import os

from pyrogram import Client, filters
from pyrogram.types import Message
from utils.misc import modules_help, prefix
from utils.scripts import format_exc, import_library

clarifai = import_library('clarifai')

from clarifai.client.model import Model


@Client.on_message(filters.command('cdxl', prefix) & filters.me)
async def cdxl(c: Client, message: Message):
    try:
        chat_id = message.chat.id
        await message.edit_text('<code>Please Wait...</code>')

        if len(message.command) > 1:
            prompt = message.text.split(maxsplit=1)[1]
        elif message.reply_to_message:
            prompt = message.reply_to_message.text
        else:
            await message.edit_text(f'<b>Usage: </b><code>{prefix}vdxl [prompt/reply to prompt]</code>')
            return

        inference_params = dict(width=1024, height=1024, steps=50, cfg_scale=9.0)

        model_prediction = Model(
            'https://clarifai.com/stability-ai/stable-diffusion-2/models/stable-diffusion-xl'
        ).predict_by_bytes(prompt.encode(), input_type='text', inference_params=inference_params)

        output_base64 = model_prediction.outputs[0].data.image.base64

        with open('sdxl_out.png', 'wb') as f:
            f.write(output_base64)

            await message.delete()
            await c.send_photo(
                chat_id,
                photo='sdxl_out.png',
                caption=f'<b>Prompt:</b><code>{prompt}</code>',
            )
            os.remove('sdxl_out.png')

    except Exception as e:
        await message.edit_text(f'An error occurred: {format_exc(e)}')


modules_help['cdxl'] = {
    'cdxl [prompt/reply to prompt]*': 'Text to Image with SDXL model',
}