you-died/you-died.py

from telegram.ext import (
    Updater,
    MessageHandler,
    Filters,
    CallbackContext,
    ConversationHandler,
    CommandHandler,
)
from telegram import Update, Message, ReplyKeyboardMarkup, ReplyKeyboardRemove, File
import json
import sys
import io
import logging
from PIL import Image

from overlay import overlay

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)


def overlay_photo(update: Update, context: CallbackContext):
    try:
        caption = update.message.caption
        if not caption:
            caption = "YOU DIED"
        file = context.bot.getFile(update.message.photo[-1].file_id)
        bytes = file.download_as_bytearray()
        image = Image.open(io.BytesIO(bytes))
        overlayed = overlay(image, caption).convert("RGB")
        output = io.BytesIO()
        overlayed.save(output, "JPEG")
        output.seek(0)
        context.bot.send_photo(chat_id=update.effective_chat.id, photo=output)
    except:
        context.bot.send_message(
            chat_id=update.effective_chat.id,
            text="Something went wrong. That might not have been an image.",
        )


def overlay_document(update: Update, context: CallbackContext):
    try:
        caption = update.message.caption
        if not caption:
            caption = "YOU DIED"
        file = context.bot.getFile(update.message.document.file_id)
        bytes = file.download_as_bytearray()
        image = Image.open(io.BytesIO(bytes))
        overlayed = overlay(image, caption)
        output = io.BytesIO()
        overlayed.save(output, "PNG")
        output.seek(0)
        context.bot.send_document(chat_id=update.effective_chat.id, document=output)
    except:
        context.bot.send_message(
            chat_id=update.effective_chat.id,
            text="Something went wrong. That might not have been an image.",
        )


def echo(update: Update, context: CallbackContext):
    context.bot.send_message(
        chat_id=update.effective_chat.id,
        text="Send an image (inline or as a file). You can customize the text by adding a caption. Alternatively, use the /caption command to choose the image, text, and color.",
    )


IMAGE, TEXT, COLOR = range(3)


def start_overlay(update: Update, context: CallbackContext) -> int:
    """Begins the overlay process."""
    update.message.reply_text("First, send the image." "\n\n" "Send /cancel to stop.")

    return IMAGE


def receive_image(update: Update, context: CallbackContext) -> int:
    """Receives the image to add text to."""
    user = update.message.from_user
    image = update.message.photo[-1].get_file()
    context.user_data["image"] = image
    logging.info("Received photo from %s", user.first_name)

    update.message.reply_text("Next, send the caption.")

    return TEXT


def receive_text(update: Update, context: CallbackContext) -> int:
    """Receives the caption to put on the image."""
    user = update.message.from_user
    text = update.message.text
    context.user_data["text"] = text
    options = [["Red", "Gold"]]

    logging.info("Received text from %s", user.first_name)

    update.message.reply_text(
        "Finally, what color do you want?",
        reply_markup=ReplyKeyboardMarkup(
            options, one_time_keyboard=True, input_field_placeholder="Color"
        ),
    )

    return COLOR


def receive_color(update: Update, context: CallbackContext) -> int:
    """Receives the color for the caption."""
    user = update.message.from_user
    color = update.message.text

    logging.info("Received color from %s", user.first_name)

    image: File = context.user_data["image"]
    text: str = context.user_data["text"]

    try:
        bytes = image.download_as_bytearray()
        image = Image.open(io.BytesIO(bytes))
        color_rgb = (200, 25, 25)

        if color == "Gold":
            color_rgb = (255, 255, 108)

        overlayed = overlay(image, text, color_rgb).convert("RGB")
        output = io.BytesIO()
        overlayed.save(output, "JPEG")
        output.seek(0)
        context.bot.send_photo(chat_id=update.effective_chat.id, photo=output)
    except:
        context.bot.send_message(
            chat_id=update.effective_chat.id,
            text="Something went wrong. That might not have been an image.",
        )

    update.message.reply_text("Done!", reply_markup=ReplyKeyboardRemove())

    return ConversationHandler.END


def cancel_overlay(update: Update, context: CallbackContext) -> int:
    """Cancels the overlay process."""
    user = update.message.from_user
    logging.info("Canceled overlay for %s", user.first_name)
    update.message.reply_text("Cancelled.", reply_markup=ReplyKeyboardRemove())

    return ConversationHandler.END


if __name__ == "__main__":
    try:
        config = json.load(open("config.json", "r", encoding="utf-8"))
    except:
        logging.error("Couldn't read the config!")
        sys.exit(1)

    updater = Updater(token=config["token"], use_context=True)
    dispatcher = updater.dispatcher
    overlay_photo_handler = MessageHandler(Filters.photo, overlay_photo)
    overlay_document_handler = MessageHandler(Filters.document, overlay_document)
    echo_handler = MessageHandler(Filters.text, echo)

    conv_handler = ConversationHandler(
        entry_points=[CommandHandler("caption", start_overlay)],
        states={
            IMAGE: [MessageHandler(Filters.photo, receive_image)],
            TEXT: [MessageHandler(Filters.text, receive_text)],
            COLOR: [MessageHandler(Filters.regex("^(Red|Gold)$"), receive_color)],
        },
        fallbacks=[CommandHandler("cancel", cancel_overlay)],
    )

    dispatcher.add_handler(conv_handler)

    dispatcher.add_handler(overlay_photo_handler)
    dispatcher.add_handler(overlay_document_handler)
    dispatcher.add_handler(echo_handler)

    updater.start_polling()