1from functools import partial
  2import os, logging
  3
  4
  5from dotenv import load_dotenv
  6from telegram import Update
  7from telegram.ext import ApplicationBuilder, CallbackQueryHandler, CommandHandler, MessageHandler, \
  8    PicklePersistence, filters, PersistenceInput, ContextTypes
  9from telegram.error import TelegramError
 10
 11from .utils import get_all, get_image, get_message_content
 12from .effects.functions import img_to_bio
 13from .slot import spin, autospin, bet, cash
 14from .effects import effectsDict
 15from .localization import get_localized_string as l, format_lang, get_lang, langs, lang_markup
 16
 17async def effect_handler(update: Update, context: ContextTypes.DEFAULT_TYPE, effect_name: str):
 18    try:
 19        effect_check, effect_fn = effectsDict[effect_name]
 20    except KeyError:
 21        raise TelegramError("effect not supported: " + effect_name)
 22    
 23    content, image, markup = await get_all(update, effect_check, context)
 24
 25    if image is None:
 26        await update.message.reply_text(l("no_caption", context))
 27        return
 28
 29    image = effect_fn(content, image)
 30
 31    if image is None:
 32        await update.message.reply_text(l("failed_effect", context))
 33        return
 34
 35    await update.message.reply_photo(photo=image, reply_markup=markup)
 36
 37async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
 38    await update.message.reply_text(l("welcome", context))
 39
 40
 41async def lewd(update: Update, context: ContextTypes.DEFAULT_TYPE):
 42    try:
 43        output = False if context.chat_data["lewd"] else True
 44    except KeyError:
 45        output = True
 46
 47    context.chat_data['lewd'] = output
 48    message = l("lewd_toggle", context).format(l("enabled", context) if output else l("disabled", context))
 49    await update.message.reply_text(message)
 50
 51
 52async def pic(update: Update, context: ContextTypes.DEFAULT_TYPE):
 53    image, markup = get_image(context)
 54    await update.message.reply_photo(photo=img_to_bio(image), parse_mode="markdown", reply_markup=markup)
 55
 56
 57async def caps(update: Update, context: ContextTypes.DEFAULT_TYPE):
 58    _, reply, _ = await get_message_content(update.message.reply_to_message, ' '.join(context.args))
 59    await update.message.reply_text(reply.upper())
 60
 61
 62async def _set_lang(update: Update, context: ContextTypes.DEFAULT_TYPE, lang: str):
 63    context.chat_data["lang"] = lang
 64    response = l("language_set", context).format(format_lang(lang))
 65    await update.message.reply_text(response)
 66
 67
 68async def lang(update: Update, context: ContextTypes.DEFAULT_TYPE):
 69    try:
 70        selected = str(context.args[0])
 71    except IndexError:
 72        selected = None
 73
 74    if selected is None:
 75        lang = format_lang(get_lang(context))
 76        choices = ", ".join(langs) + "."
 77        text = l("current_language", context).format(lang, choices)
 78        return await update.message.reply_text(text,reply_markup=lang_markup)
 79
 80    if selected not in langs:
 81        return await update.message.reply_text(text=l("invalid_language", context))
 82
 83    return await _set_lang(update, context, selected)
 84
 85
 86def unknown(update: Update, context: ContextTypes.DEFAULT_TYPE):
 87    logging.info(f"User {update.message.from_user.full_name} sent {update.message.text_markdown_v2} and I don't know what that means.")
 88
 89
 90async def error_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
 91    try:
 92        raise context.error
 93    except Exception as e:
 94        logging.error(str(e))
 95        await update.message.reply_text(l('error', context))
 96
 97
 98def _add_effect_handler(application: ApplicationBuilder, command: str):
 99    callback = partial(effect_handler, effect_name=command)
100
101    application.add_handler(CommandHandler(command, callback))
102    application.add_handler(MessageHandler(filters.Caption([f"/{command}"]), callback))
103
104async def keyboard_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
105    query = update.callback_query
106    data = query.data
107
108    if data.startswith("reroll"):
109        amount = int(data.split(" ")[1])
110
111        if amount <= 1:
112            return await spin(update, context)
113        return await autospin(context, update.effective_chat.id, amount)
114
115    match data:
116        case "none":
117            return query.answer(l("none_callback", context))
118        case "set_lang_en":
119            lang = "en"
120            await _set_lang(update, context, lang)
121            return await query.answer(l("language_set", context).format(format_lang(lang)))
122        case "set_lang_it":
123            lang = "it"
124            await _set_lang(update, context, lang)
125            return await query.answer(l("language_set", context).format(format_lang(lang)))
126        case other:
127            logging.error(f"unknown callback: {data}")
128
129    return await query.answer()
130
131
132def main():
133    load_dotenv()
134    logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
135
136    token = os.getenv("token")
137    pers = PersistenceInput(bot_data=False, callback_data=False)
138    persistence = PicklePersistence(filepath='bot-data.pkl', store_data=pers)
139
140    application = ApplicationBuilder().token(token).persistence(persistence).build()
141
142    application.add_error_handler(error_callback)
143    application.add_handler(CallbackQueryHandler(callback=keyboard_handler))
144
145    # commands
146    application.add_handler(CommandHandler('start', start))
147    application.add_handler(CommandHandler('lang', lang))
148    application.add_handler(CommandHandler('lewd', lewd))
149    application.add_handler(CommandHandler('caps', caps))
150    application.add_handler(CommandHandler('pic', pic))
151
152    # effects
153    _add_effect_handler(application, 'ttbt')
154    _add_effect_handler(application, 'tt')
155    _add_effect_handler(application, 'bt')
156    _add_effect_handler(application, 'splash')
157    _add_effect_handler(application, 'wot')
158    _add_effect_handler(application, 'text')
159
160    # games
161    application.add_handler(CommandHandler('spin', spin))
162    application.add_handler(CommandHandler('bet', bet))
163    application.add_handler(CommandHandler('cash', cash))
164
165    # fallback
166    application.add_handler(MessageHandler(filters.COMMAND, unknown))
167    application.run_polling()
168
169
170if __name__ == "__main__":
171    main()