import runpod
import os
import torch
import io
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline
import os
import telegram


BOT_TOKEN = os.environ.get("BOT_TOKEN")
bot = telegram.Bot(token=BOT_TOKEN)

device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "logo-marks-6-1_5", torch_dtype=torch.float16, safety_checker=None
).to(device)


def generate_image(user_id, seed, image_order, prompt=""):

    init_image = Image.open(
        f"logo-marks-6-1_5/reference-images/{image_order}.png"
    ).convert("RGB")
    init_image.thumbnail((768, 768))

    folder_path = f"images/{user_id}"
    os.makedirs(folder_path, exist_ok=True)

    generator = torch.Generator(device="cuda").manual_seed(seed)
    images = pipe(
        prompt=f"{prompt} logo mark in zvxsbj style",
        image=init_image,
        strength=0.95,
        guidance_scale=30,
        generator=generator,
    ).images
    print(f"Image with seed {seed} and image_order {image_order} saved")
    return images[0]


def send_ranom_ask(chat_id):
    bot.send_message(
        chat_id=chat_id,
        text=f"Press /random to get random logo or send a short prompt (1-2 word) and i will generate 10 logotypes for you",
        reply_markup=telegram.ReplyKeyboardMarkup(
            [
                ["/random"],
            ],
            resize_keyboard=True,
            one_time_keyboard=True,
        ),
    )


def handler(event):
    user_id = event["input"]["user_id"]
    prompt = event["input"]["prompt"]

    image_requests = event["input"]["image_requests"]  # example: [ [1, 0], [1, 1] ]

    image_bytes = []

    for image_request in image_requests:
        image = generate_image(user_id, image_request[0], image_request[1], prompt)
        bytes_image = io.BytesIO()
        image.save(bytes_image, "PNG")
        bytes_image.seek(0)
        image_bytes.append(telegram.InputMediaPhoto(bytes_image))

    bot.send_media_group(user_id, image_bytes)

    send_ranom_ask(user_id)

    return "Hello World"


runpod.serverless.start({"handler": handler})