from functools import partial
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Any

import gradio as gr
import segno
from gradio.components import Component
from huggingface_hub import InferenceClient
from PIL import Image
from qrcode_artistic import write_artistic

try:
    import dotenv

    dotenv.load_dotenv()
except ImportError:
    pass

client = InferenceClient(model="black-forest-labs/FLUX.1-schnell")

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(
                "https://wheelingvultures.bandcamp.com/album/ep", label="Text"
            )
            prompt = gr.TextArea("A psychedelic vulture", label="Prompt")
            button = gr.Button("Generate")
        with gr.Column():
            background = gr.Image(visible=False, type="filepath")
            scale = gr.Slider(3, 15, 9, step=1, label="Scale")
            output = gr.Image()

    gr.on(
        [button.click, prompt.submit],
        partial(gr.update, interactive=False),
        outputs=button,
    ).then(
        partial(client.text_to_image, width=400, height=400),
        inputs=prompt,
        outputs=background,
    ).then(
        partial(gr.update, interactive=True),
        outputs=button,
    )

    @gr.on(
        triggers=[text.submit, background.change, scale.change],
        inputs={text, background, scale},
        outputs=output,
    )
    def generate_code(data: dict[Component, Any]):
        if data.get(background) is None:
            return None

        qr_code = segno.make(data[text], error="h")
        suffix = Path(data[background]).suffix

        with NamedTemporaryFile(suffix=suffix) as temp_file:
            write_artistic(
                qr_code,
                target=temp_file,
                background=data[background],
                scale=data[scale],
            )

            return Image.open(temp_file.name)


if __name__ == "__main__":
    demo.launch()