import os
import random
import time
from typing import Any

import gradio as gr
import pillow_avif  # noqa: F401
import pillow_heif
import spaces
import torch
from gradio_imageslider import ImageSlider
from PIL import Image


class ModuleInterface:
    def __init__(self, d: dict[str, Any]):
        self.d = d

    @property
    def Model(self) -> Any:
        return self.d["Model"]

    def process(self, input_image: Image.Image, model: Any, seed: int) -> Image.Image:
        return self.d["process"](input_image, model, seed)


assert (src := os.getenv("LIGHT_SWITCHER_LITE")), "LIGHT_SWITCHER_LITE not set"
exec_globals: dict[str, Any] = {}
exec(src, exec_globals)
light_switcher_lite = ModuleInterface(exec_globals)

pillow_heif.register_avif_opener()


DEVICE_CPU = torch.device("cpu")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32

# CPU -> GPU dance because of ZeroGPU

path = "finegrain/weights-light-switcher-space"
model = light_switcher_lite.Model.from_pretrained(path, device=DEVICE_CPU, dtype=DTYPE)

model.to(DEVICE)


@spaces.GPU
def process(input_image: Image.Image, seed: int = 42) -> tuple[tuple[Image.Image, Image.Image], dict[str, Any]]:
    output_image = light_switcher_lite.process(input_image, model, seed)
    resized_input_image = input_image.resize(output_image.size)
    return ((resized_input_image, output_image), gr.update(value=random.choice(BUTTON_LABELS)))


TITLE = """
<h1>Finegrain Light Switcher (Lite Version)</h1>
<p>
  Given an image with a lamp switched off, the model should turn it on.
</p>
<p>
  🔌 For higher resolution results with control over lighting intensity and warmth,&nbsp;
  <a href="https://finegrain.ai">head to the Finegrain API</a> 🔌
</p>
<p>
  <a href="https://discord.gg/zFKg5TjXub" target="_blank">[Discord]</a>
  <a href="https://github.com/finegrain-ai" target="_blank">[GitHub]</a>
  <a href="https://finegrain.ai">[Finegrain API]</a>
</p>
"""

BUTTON_LABELS = [
    "💡",
    "Let there be light!",
    "Light it up like a Christmas tree!",
    "Turn it on!",
    "Aziz, Light!",
    "Make it shine ✨",
    "Flip the magic switch.",
]

random.seed(time.time())

with gr.Blocks() as demo:
    gr.HTML(TITLE)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.ClearButton(components=None, value=random.choice(BUTTON_LABELS))
        with gr.Column():
            output_slider = ImageSlider(label="Before / After")
            run_button.add(output_slider)

    with gr.Accordion("Advanced Options", open=False):
        seed = gr.Slider(minimum=0, maximum=999, value=42, step=1, label="Seed")

    run_button.click(
        fn=process,
        inputs=[input_image, seed],
        outputs=[output_slider, run_button],
    )

    gr.Examples(
        examples=[
            "examples/01.webp",
            "examples/02.webp",
            "examples/03.webp",
            "examples/04.webp",
            "examples/05.webp",
        ],
        inputs=[input_image],
        outputs=[output_slider, run_button],
        fn=process,
        cache_examples=True,
        run_on_click=False,
    )

demo.launch(share=False)