from pathlib import Path
from tempfile import NamedTemporaryFile

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

if gr.NO_RELOAD:
    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 vulture in black and white", 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(
        triggers=[button.click, prompt.submit],
        inputs=prompt,
        outputs=[background, output],
    )
    def generate_background(prompt):
        return client.text_to_image(prompt, width=400, height=400), None

    @gr.on(
        triggers=[text.submit, scale.change, background.change],
        inputs=[text, scale, background],
        outputs=output,
    )
    def generate_code(text, scale, background):
        if background is None:
            return None

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

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

            return Image.open(temp_file.name)


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