import os
import json
import time
import math
import tempfile
import uuid

import spaces
import torch
import gradio as gr
from PIL import Image

from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image

MODEL_ID = "stabilityai/sdxl-turbo"
TORCH_DTYPE = torch.float16
SESSION_DURATION = 58  # seconds; leave 2s slack on @spaces.GPU(duration=60)
POLL_INTERVAL = 0.05  # seconds between slot checks in the streaming loop

SESSION_DIR = os.path.join(tempfile.gettempdir(), "sdxl_turbo_sessions")
os.makedirs(SESSION_DIR, exist_ok=True)


# ZeroGPU snapshot preload at module scope (import spaces hijacked CUDA above).
i2i_pipe = AutoPipelineForImage2Image.from_pretrained(
    MODEL_ID, safety_checker=None, torch_dtype=TORCH_DTYPE, variant="fp16"
).to("cuda")
i2i_pipe.set_progress_bar_config(disable=True)

t2i_pipe = AutoPipelineForText2Image.from_pretrained(
    MODEL_ID, safety_checker=None, torch_dtype=TORCH_DTYPE, variant="fp16"
).to("cuda")
t2i_pipe.set_progress_bar_config(disable=True)


# ---------- session-slot IPC ----------
# The @spaces.GPU streaming generator runs in a forked process and can't see
# parent-process state. We pass live inputs via an atomically-replaced JSON
# file keyed by session_hash. Change handlers (no GPU) write to it; the
# generator polls it. Atomic write via os.replace avoids torn reads.

def _slot_path(session_id: str) -> str:
    return os.path.join(SESSION_DIR, f"{session_id}.json")


def _read_slot(session_id: str):
    try:
        with open(_slot_path(session_id)) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return None


def _write_slot(session_id: str, data: dict) -> None:
    p = _slot_path(session_id)
    tmp = p + f".{uuid.uuid4().hex}.tmp"
    with open(tmp, "w") as f:
        json.dump(data, f)
    os.replace(tmp, p)


def _bump_slot(session_id: str, **fields) -> None:
    slot = _read_slot(session_id) or {"version": 0}
    slot.update(fields)
    slot["version"] = slot.get("version", 0) + 1
    _write_slot(session_id, slot)


def update_slot(prompt, strength, steps, seed, image_path, req: gr.Request):
    _bump_slot(
        req.session_hash,
        prompt=prompt or "",
        strength=float(strength),
        steps=int(steps),
        seed=int(seed),
        image_path=image_path,
    )


# ---------- inference ----------

def resize_crop(image: Image.Image, size: int = 512) -> Image.Image:
    image = image.convert("RGB")
    w, h = image.size
    return image.resize((size, int(size * (h / w))), Image.BICUBIC)


def _generate(slot: dict):
    prompt = slot.get("prompt", "")
    if not prompt:
        return None
    seed = int(slot.get("seed", 0))
    steps = int(slot.get("steps", 2))
    strength = float(slot.get("strength", 0.7))
    image_path = slot.get("image_path")

    generator = torch.manual_seed(seed)
    if image_path:
        init_image = resize_crop(Image.open(image_path))
        if int(steps * strength) < 1:
            steps = math.ceil(1 / max(0.10, strength))
        result = i2i_pipe(
            prompt=prompt,
            image=init_image,
            generator=generator,
            num_inference_steps=steps,
            guidance_scale=0.0,
            strength=strength,
            width=512,
            height=512,
            output_type="pil",
        )
    else:
        result = t2i_pipe(
            prompt=prompt,
            generator=generator,
            num_inference_steps=max(1, steps),
            guidance_scale=0.0,
            width=512,
            height=512,
            output_type="pil",
        )
    return result.images[0]


@spaces.GPU(duration=60)
def run_session(prompt, strength, steps, seed, image_path, req: gr.Request):
    """Streaming generator: holds the GPU for the session window and yields
    a new image whenever any input changes."""
    sid = req.session_hash
    _write_slot(sid, {
        "prompt": prompt or "",
        "strength": float(strength),
        "steps": int(steps),
        "seed": int(seed),
        "image_path": image_path,
        "version": 0,
    })

    deadline = time.time() + SESSION_DURATION
    last_version = -1
    last_image = None
    while time.time() < deadline:
        slot = _read_slot(sid)
        if slot is not None and slot.get("version", 0) != last_version:
            last_version = slot["version"]
            try:
                img = _generate(slot)
            except Exception as e:
                print(f"[session {sid[:8]}] generate err: {e}")
                img = None
            if img is not None:
                last_image = img
                yield img
        else:
            time.sleep(POLL_INTERVAL)

    if last_image is not None:
        yield last_image


def _timer_html():
    return f"""
    <div id="sdxl-timer" style="font-family:monospace;font-size:1.1em;padding:8px 12px;background:#222;color:#9f9;border-radius:8px;display:inline-block;">
      ⏱ Session active — <span id="sdxl-countdown">{SESSION_DURATION}</span>s left
    </div>
    <script>
    (function() {{
      const el = document.getElementById('sdxl-countdown');
      if (!el) return;
      let remaining = {SESSION_DURATION};
      const t = setInterval(() => {{
        remaining -= 1;
        if (el) el.textContent = Math.max(0, remaining);
        if (remaining <= 0) clearInterval(t);
      }}, 1000);
    }})();
    </script>
    """


def start_session_ui():
    return (
        gr.update(visible=False),       # start_btn hidden
        gr.update(interactive=True),    # prompt enabled
        gr.update(interactive=True),    # generate_bt enabled
        gr.update(value=_timer_html(), visible=True),  # countdown shown
    )


def end_session_ui(req: gr.Request):
    try:
        os.remove(_slot_path(req.session_hash))
    except FileNotFoundError:
        pass
    return (
        gr.update(visible=True),        # start_btn back
        gr.update(interactive=False),   # prompt disabled
        gr.update(interactive=False),   # generate_bt disabled
        gr.update(value="", visible=False),  # countdown hidden
    )


# ---------- UI ----------

css = """
#container { margin: 0 auto; max-width: 80rem; }
#intro { max-width: 100%; text-align: center; margin: 0 auto; }
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="container"):
        gr.Markdown(
            """# SDXL Turbo • Realtime ZeroGPU Session
            ## Unofficial Demo
            Click **Start Realtime Session** to acquire ZeroGPU for 60 seconds.
            While the session is active, every change to the prompt, sliders or
            input image updates the output in realtime — no per-keystroke GPU
            fork.
            **Model**: https://huggingface.co/stabilityai/sdxl-turbo
            """,
            elem_id="intro",
        )

        start_btn = gr.Button("▶️  Start Realtime Session (60s)", variant="primary")
        timer = gr.HTML("", visible=False)

        with gr.Row():
            prompt = gr.Textbox(
                placeholder="Insert your prompt here:",
                scale=5,
                container=False,
                interactive=False,
            )
            generate_bt = gr.Button("Generate", scale=1, interactive=False)

        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    sources=["upload", "webcam", "clipboard"],
                    label="Webcam / Upload (optional — leaves text-to-image mode if empty)",
                    type="filepath",
                )
            with gr.Column():
                image = gr.Image(type="pil", label="Output")
                with gr.Accordion("Advanced options", open=False):
                    strength = gr.Slider(label="Strength", value=0.7, minimum=0.0, maximum=1.0, step=0.001)
                    steps = gr.Slider(label="Steps", value=2, minimum=1, maximum=10, step=1)
                    seed = gr.Slider(randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1)

        inputs = [prompt, strength, steps, seed, image_input]

        start_btn.click(
            fn=start_session_ui,
            inputs=None,
            outputs=[start_btn, prompt, generate_bt, timer],
            queue=False,
        ).then(
            fn=run_session,
            inputs=inputs,
            outputs=image,
            show_progress=False,
        ).then(
            fn=end_session_ui,
            inputs=None,
            outputs=[start_btn, prompt, generate_bt, timer],
            queue=False,
        )

        # All change handlers update the session slot WITHOUT touching the GPU.
        for component in [prompt, strength, steps, seed, image_input]:
            component.change(
                fn=update_slot,
                inputs=inputs,
                outputs=None,
                queue=False,
                show_progress=False,
            )
        # Generate button just bumps slot version with current inputs
        generate_bt.click(
            fn=update_slot,
            inputs=inputs,
            outputs=None,
            queue=False,
            show_progress=False,
        )

demo.queue()
demo.launch()