import re
from functools import partial
from io import BytesIO
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")

MODELS = [
    "stabilityai/stable-diffusion-3.5-large",
    "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")
            model = gr.Radio(MODELS, value=MODELS[0], label="Model")
            button = gr.Button("Generate")

        with gr.Column():
            output = gr.Image()
            background = gr.Image(visible=False, type="filepath")
            scale = gr.Slider(3, 15, 9, step=1, label="Scale")

            with gr.Row():
                color_dark = gr.ColorPicker("#000000", label="Dark")
                color_light = gr.ColorPicker("#FFFFFF", label="Light")

    def generate_background(data: dict[Component, Any]):
        if not data.get(prompt):
            return gr.skip(), gr.skip()

        return client.text_to_image(data[prompt], model=data[model]), None

    def generate_output(data: dict[Component, Any]):
        if data.get(background) is None:
            return None

        def to_hex_format(value: str):
            if value.startswith("#"):
                return value

            matches = re.findall(r"\d+(?:\.\d+)?", value)
            r, g, b = map(int, map(float, matches[:3]))

            return f"#{r:x}{g:x}{b:x}".upper()

        image = Image.open(data[background])
        qr_code = segno.make(data[text], error="h")

        with BytesIO() as buffer:
            write_artistic(
                qr_code,
                target=buffer,
                background=image.fp,
                kind=image.format,
                scale=data[scale],
                light=to_hex_format(data[color_light]),
                dark=to_hex_format(data[color_dark]),
            )

            return Image.open(buffer)

    gr.on(
        [button.click, prompt.submit],
        partial(gr.update, interactive=False),
        outputs=button,
    ).then(
        generate_background,
        inputs={prompt, model},
        outputs=[background, output],
    ).then(
        partial(gr.update, interactive=True),
        outputs=button,
    )

    gr.on(
        [
            text.submit,
            background.change,
            scale.change,
            color_light.change,
            color_dark.change,
        ],
        generate_output,
        inputs={
            text,
            background,
            scale,
            color_light,
            color_dark,
        },
        outputs=output,
    )


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