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, InlineKeyboardMarkup, InlineKeyboardButton
 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, image=None):
 35    if text.strip() == "":
 36        return None, l("no_caption", lang)
 37    
 38    markup = ""
 39    if image is None:
 40        image, markup = _get_image(context, bio=False)
 41    image = tt_bt_effect(text, image)
 42    return _img_to_bio(image), markup
 43
 44def _get_reply(input, context, fallback=""):
 45
 46    if input is None:
 47        return None, fallback
 48    
 49    image = None
 50    if len(input.photo) > 0:
 51        image = input.photo[-1].get_file()
 52        #image = context.bot.get_file(input.photo[-1].file_id)
 53        image = Image.open(BytesIO(image.download_as_bytearray()))
 54    
 55    if input.caption is not None:
 56        return image, input.caption
 57    
 58    if input.text is not None:
 59        return image, input.text
 60    
 61    return image, fallback
 62
 63def _get_lewd(context):
 64    try:
 65        return context.chat_data["lewd"]
 66    except KeyError:
 67        return False
 68
 69def _get_image(context, tag="", bio=True):
 70    if context is not None:
 71        if _get_lewd(context):
 72            image, url = get_random_image(rating_lewd, tag)
 73        else:
 74            image, url = get_random_image(rating_normal, tag)
 75    
 76    if image is None:
 77        logging.warning("Getting Image failed")
 78        raise TelegramError("bad image")
 79    
 80    markup = InlineKeyboardMarkup([[InlineKeyboardButton(text=l("sauce", lang), url=url)]])
 81    
 82    if bio:
 83        return _img_to_bio(image), markup
 84    return image, markup
 85    
 86
 87def start(update: Update, context: CallbackContext):
 88    context.bot.send_message(chat_id=update.effective_chat.id, text=l("welcome", lang))
 89
 90def set_lewd(update: Update, context: CallbackContext):
 91    try:
 92        output = False if context.chat_data["lewd"] else True
 93    except KeyError:
 94        output = True
 95        
 96    context.chat_data['lewd'] = output
 97    message = l("lewd_toggle", lang).format(l("enabled", lang) if output else l("disabled", lang))
 98    
 99    context.bot.send_message(chat_id=update.effective_chat.id, text=message)
100
101def pic(update: Update, context: CallbackContext):
102    try:
103        tag = " " + context.args[0]
104    except IndexError:
105        tag = ""
106
107    image, markup = _get_image(context, tag)
108    update.message.reply_photo(photo=image, parse_mode="markdown", reply_markup=markup)
109
110def raw(update: Update, context: CallbackContext):
111    tag = ""
112    try:
113        tag += " " + context.args[0]
114        tag += " " + context.args[1]
115    except IndexError:
116        pass
117    
118    image, url = get_random_image("", tag)
119    
120    if image is None:
121        logging.warning("Getting Image failed")
122        raise TelegramError("bad image")
123    
124    image = _img_to_bio(image)
125    markup = InlineKeyboardMarkup([[InlineKeyboardButton(text=l("sauce", lang), url=url)]])
126    
127    update.message.reply_photo(photo=image, parse_mode="markdown", reply_markup=markup)
128
129def pilu(update: Update, context: CallbackContext):
130    logging.warning(f"User {update.message.from_user.username} requested an explicit pic.")
131    try:
132        tag = " " + context.args[0]
133    except IndexError:
134        tag = ""
135    image, url = get_random_image("rating:e" + tag)
136    if image is None:
137        logging.warning("Getting Image failed")
138        raise TelegramError("bad image")
139        return
140    
141    image = _img_to_bio(image)
142    markup = InlineKeyboardMarkup([[InlineKeyboardButton(text=l("sauce", lang), url=url)]])
143    
144    update.message.reply_photo(photo=image, caption=url, parse_mode="markdown", reply_markup=markup)
145
146def tt(update: Update, context: CallbackContext):
147    image, reply = _get_reply(update.message.reply_to_message, context)
148    content = _get_args(context)
149    input_text = f"{reply} {content}".replace("\n", " ")
150    
151    image, markup = _ttbt_general(context, input_text, image)
152    
153    if image is None:
154        update.message.reply_text(markup)
155        return
156    update.message.reply_photo(photo=image, parse_mode="markdown", reply_markup=markup)
157
158def bt(update: Update, context: CallbackContext):
159    image, reply = _get_reply(update.message.reply_to_message, context)
160    content = _get_args(context)
161    input_text = f"{reply} {content}".replace("\n", " ")
162    
163    image, markup =_ttbt_general(context, " \n" + input_text, image)
164    
165    if image is None:
166        update.message.reply_text(markup)
167        return
168    update.message.reply_photo(photo=image, parse_mode="markdown", reply_markup=markup)
169
170def ttbt(update: Update, context: CallbackContext):
171    message = update.message
172    
173    image, reply = _get_reply(message.reply_to_message, context)
174    content = message.text.split(" ")
175    content.pop(0)
176    content = " ".join(content)
177    
178    input_text = f"{reply}\n{content}"
179    
180    image, markup =_ttbt_general(context, input_text, image)
181    
182    if image is None:
183        message.reply_text(markup)
184        return
185    message.reply_photo(photo=image, parse_mode="markdown", reply_markup=markup)
186    
187def caps(update: Update, context: CallbackContext):
188    _, reply = _get_reply(update.message.reply_to_message, context, _get_args(context))
189    context.bot.send_message(chat_id=update.effective_chat.id, text=reply.upper())
190    
191def unknown(update: Update, context: CallbackContext):
192    logging.info(f"User {update.message.from_user.username} sent {update.message.text_markdown_v2} and I don't know what that means.")
193    
194def error_callback(update: Update, context: CallbackContext):
195    try:
196        raise context.error
197    #except BadRequest:
198    #    logging.error("BadRequest!!")
199    except TelegramError:
200        logging.error("TelegramError!!")
201        context.bot.send_message(chat_id=update.effective_chat.id, text=l('error', lang))
202def main():
203    
204    updater = Updater(token=os.getenv("token"))
205    dispatcher = updater.dispatcher
206    dispatcher.add_error_handler(error_callback)
207    
208    dispatcher.add_handler(CommandHandler('start', start))
209    dispatcher.add_handler(CommandHandler('lewd', set_lewd))
210    dispatcher.add_handler(CommandHandler('caps', caps))
211    dispatcher.add_handler(CommandHandler('pic', pic))
212    dispatcher.add_handler(CommandHandler('raw', raw))
213    dispatcher.add_handler(CommandHandler('pilu', pilu))
214    dispatcher.add_handler(CommandHandler('ttbt', ttbt))
215    dispatcher.add_handler(CommandHandler('tt', tt))
216    dispatcher.add_handler(CommandHandler('bt', bt))
217    
218    dispatcher.add_handler(MessageHandler(Filters.command, unknown))
219    updater.start_polling()
220    updater.idle()
221
222if __name__ ==  "__main__":
223    main()