1from PIL import Image
  2from Api import get_random_image, rating_normal, rating_lewd
  3from Effects import tt_bt_effect
  4import Constants as C
  5from Constants import get_localized_string as l
  6
  7from dotenv import load_dotenv
  8load_dotenv()
  9import os, logging
 10logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
 11from io import BytesIO
 12
 13from telegram.error import TelegramError, BadRequest
 14from telegram.ext import Updater, CallbackContext, CommandHandler, MessageHandler, Filters
 15from telegram import Update
 16
 17lang = 'us'
 18
 19def _get_args(context):
 20    logging.info(context.args)
 21    return ' '.join(context.args)
 22
 23def _img_to_bio(image):
 24    bio = BytesIO()
 25    #bio.name = 'image.jpeg'
 26    
 27    if image.mode in ("RGBA", "P"):
 28        image = image.convert("RGB")
 29
 30    image.save(bio, 'JPEG')
 31    bio.seek(0)
 32    return bio
 33
 34def _ttbt_general(context, text):
 35    if text.strip() == "":
 36        return None, l("no_caption", lang)
 37    
 38    image, url = _get_image(context, bio=False)
 39    image = tt_bt_effect(text, image)
 40    return _img_to_bio(image), url
 41
 42def _get_reply(input, fallback=""):
 43    if input is None:
 44        return fallback
 45    return input.text
 46
 47def _get_lewd(context):
 48    try:
 49        return context.chat_data["lewd"]
 50    except KeyError:
 51        return False
 52
 53def _get_image(context, bio=True):
 54    if context is not None:
 55        if _get_lewd(context):
 56            image, url = get_random_image(rating_lewd)
 57        else:
 58            image, url = get_random_image(rating_normal)
 59    
 60    if image is None:
 61        logging.warning("Getting Image failed")
 62        raise TelegramError("bad image")
 63    
 64    url = l("sauce", lang).format(url)
 65    if bio:
 66        return _img_to_bio(image), url
 67    return image, url
 68    
 69
 70def start(update: Update, context: CallbackContext):
 71    context.bot.send_message(chat_id=update.effective_chat.id, text=l("welcome", lang))
 72
 73def set_lewd(update: Update, context: CallbackContext):
 74    try:
 75        output = False if context.chat_data["lewd"] else True
 76    except KeyError:
 77        output = True
 78        
 79    context.chat_data['lewd'] = output
 80    message = l("lewd_toggle", lang).format(l("enabled", lang) if output else l("disabled", lang))
 81    
 82    context.bot.send_message(chat_id=update.effective_chat.id, text=message)
 83
 84def pic(update: Update, context: CallbackContext):
 85    image, url = _get_image(context)
 86    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown")
 87
 88def pilu(update: Update, context: CallbackContext):
 89    logging.warning(f"User {update.message.from_user.username} requested an explicit pic.")
 90    try:
 91        tag = " " + context.args[0]
 92    except IndexError:
 93        tag = ""
 94    image, url = get_random_image("e" + tag)
 95    if image is None:
 96        logging.warning("Getting Image failed")
 97        raise TelegramError("bad image")
 98        return
 99    url = l("sauce", lang).format(url)
100    image = _img_to_bio(image)
101    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown")
102
103def tt(update: Update, context: CallbackContext):
104    reply = _get_reply(update.message.reply_to_message)
105    content = _get_args(context)
106    input_text = f"{reply} {content}".replace("\n", " ")
107    
108    image, url = _ttbt_general(context, input_text)
109    
110    if image is None:
111        update.message.reply_text(url)
112        return
113    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown")
114
115def bt(update: Update, context: CallbackContext):
116    reply = _get_reply(update.message.reply_to_message)
117    content = _get_args(context)
118    input_text = f"{reply} {content}".replace("\n", " ")
119    
120    image, url =_ttbt_general(context, " \n" + input_text)
121    
122    if image is None:
123        update.message.reply_text(url)
124        return
125    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown")
126
127def ttbt(update: Update, context: CallbackContext):
128    reply = _get_reply(update.message.reply_to_message)
129    content = _get_args(context)
130    input_text = f"{reply}\n{content}"
131    
132    image, url =_ttbt_general(context, input_text)
133    
134    if image is None:
135        update.message.reply_text(url)
136        return
137    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown")
138    
139def caps(update: Update, context: CallbackContext):
140    reply = _get_reply(update.message.reply_to_message, _get_args(context))
141    context.bot.send_message(chat_id=update.effective_chat.id, text=reply.upper())
142    
143def unknown(update: Update, context: CallbackContext):
144    logging.info(f"User {update.message.from_user.username} sent {_get_args(context)}")
145    context.bot.reply_text(text=l("unknown", lang))
146    
147def error_callback(update: Update, context: CallbackContext):
148    try:
149        raise context.error
150    #except BadRequest:
151    #    logging.error("BadRequest!!")
152    except TelegramError:
153        logging.error("TelegramError!!")
154        context.bot.send_message(chat_id=update.effective_chat.id, text=l('error', lang))
155def main():
156    
157    updater = Updater(token=os.getenv("token"))
158    dispatcher = updater.dispatcher
159    dispatcher.add_error_handler(error_callback)
160    
161    dispatcher.add_handler(CommandHandler('start', start))
162    dispatcher.add_handler(CommandHandler('lewd', set_lewd))
163    dispatcher.add_handler(CommandHandler('caps', caps))
164    dispatcher.add_handler(CommandHandler('pic', pic))
165    dispatcher.add_handler(CommandHandler('pilu', pilu))
166    dispatcher.add_handler(CommandHandler('ttbt', ttbt))
167    dispatcher.add_handler(CommandHandler('tt', tt))
168    dispatcher.add_handler(CommandHandler('bt', bt))
169    
170    dispatcher.add_handler(MessageHandler(Filters.command, unknown))
171    updater.start_polling()
172    updater.idle()
173
174if __name__ ==  "__main__":
175    main()