"""
Anima RDBT — Gradio front-end for ZeroGPU Hugging Face Spaces, with an in-process
Diffusers-based Anima pipeline (diffusers + diffusers-anima) for RDBT sampling.
"""
from __future__ import annotations

import functools
import os
import traceback
from typing import Any

import gradio as gr

# ZeroGPU: https://huggingface.co/docs/hub/main/spaces-zerogpu
try:
    import spaces
except ImportError:  # local / tests without the runtime wheel
    class _Spaces:
        @staticmethod
        def GPU(*args: Any, **kwargs: Any):
            if args and callable(args[0]) and not kwargs:
                return args[0]

            def _deco(fn: Any) -> Any:
                return fn

            return _deco

    spaces = _Spaces()  # type: ignore[misc, assignment]

from src.diffusers_backend import run_generation
from src.errors import UserFacingError
from src.validation import validate_and_clamp

from src import config

LICENSE_BANNER = (
    "**License:** Anima is distributed under the **CircleStone Labs Non-Commercial License** "
    "and is a derivative of **Cosmos-Predict2-2B-Text2Image** (NVIDIA terms where applicable). "
    "Use **non-commercially** unless you have a separate license. See the "
    "[Anima model card](https://huggingface.co/circlestone-labs/Anima)."
)

FALLBACK_NOTE = (
    "If a sampler/scheduler is not supported by the Diffusers (Anima) scheduler, the Space "
    "remaps to a supported pair (typically **euler_ancestral**-style and **simple**) and may note it in the status text."
)


def _compute_gpu_duration(
    _prompt: str,
    _neg: str,
    _w: int,
    _h: int,
    steps: float,
    _cfg: float,
    batch: float,
    _s: str,
    _sc: str,
    _d: float,
) -> int:
    try:
        s = int(steps)
    except (TypeError, ValueError):
        s = 16
    try:
        b = int(batch)
    except (TypeError, ValueError):
        b = 1
    s = max(1, min(50, s))
    b = max(1, min(config.MAX_BATCH, b))
    # Queue-friendly: high-res + batch → more time; cap 300s (ZeroGPU allowlist)
    est = 25 + s * 4 * b
    w = int(os.environ.get("ANIMA_MAX_GPU_DURATION", "300"))
    return min(int(w), max(30, int(est)))


@spaces.GPU(duration=_compute_gpu_duration)  # type: ignore[misc]
def _generate_task(
    prompt: str,
    negative_prompt: str,
    width: float,
    height: float,
    steps: float,
    cfg: float,
    batch_size: float,
    sampler_name: str,
    scheduler: str,
    denoise: float,
) -> tuple[Any, str]:
    """Returns (gallery data or None, status markdown)."""
    try:
        params = validate_and_clamp(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=steps,
            cfg=cfg,
            batch_size=batch_size,
            sampler_name=sampler_name,
            scheduler=scheduler,
            denoise=denoise,
        )
    except UserFacingError as e:
        lines = [f"**Error:** {e.user_message}"]
        if e.details:
            lines.append(f"`{e.details}`")
        return None, "\n\n".join(lines)
    try:
        images, det = run_generation(params)
    except UserFacingError as e:
        return None, f"**Error:** {e.user_message}"
    except Exception as e:  # noqa: BLE001
        tb = traceback.format_exc()
        return None, f"**Error:** Generation failed: `{e!s}`\n\n```\n{tb[-4000:]}\n```"
    status_parts = [
        "**Done.** " + det,
    ]
    if params.warnings:
        status_parts.append("**Notices:**\n" + "\n".join(f"- {w}" for w in params.warnings))
    if images:
        return images, "\n\n".join(status_parts)
    return None, "\n\n".join(status_parts)


def build_ui() -> gr.Blocks:
    # Gradio 6.0: theme and related UI options belong on launch(), not Blocks().
    with gr.Blocks(
        title="Anima RDBT (Gradio)",
    ) as demo:
        gr.Markdown(
            f"# Anima RDBT — {config.RDBT_UNET_NAME}\n\n{LICENSE_BANNER}\n\n{FALLBACK_NOTE}"
        )
        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="prompt",
                    lines=3,
                    value="digital anime illustration, 1girl, smile",
                )
                neg = gr.Textbox(
                    label="negative_prompt",
                    lines=2,
                    value="",
                )
                with gr.Row():
                    w = gr.Slider(
                        config.MIN_WH,
                        config.MAX_WH,
                        value=config.DEFAULT_WIDTH,
                        step=64,
                        label="width",
                    )
                    h = gr.Slider(
                        config.MIN_WH,
                        config.MAX_WH,
                        value=config.DEFAULT_HEIGHT,
                        step=64,
                        label="height",
                    )
                st = gr.Slider(
                    config.MIN_STEPS,
                    config.MAX_STEPS,
                    value=config.DEFAULT_STEPS,
                    step=1,
                    label="steps",
                )
                cfg = gr.Slider(
                    config.MIN_CFG,
                    config.MAX_CFG,
                    value=config.DEFAULT_CFG,
                    step=config.CFG_STEP,
                    label="cfg",
                )
            with gr.Column(scale=1):
                with gr.Accordion("Advanced", open=True):
                    batch = gr.Slider(
                        config.MIN_BATCH,
                        config.MAX_BATCH,
                        value=config.DEFAULT_BATCH_SIZE,
                        step=1,
                        label="batch_size",
                    )
                    sampler = gr.Dropdown(
                        choices=list(config.SAMPLER_CHOICES),
                        value=config.DEFAULT_SAMPLER,
                        label="sampler_name",
                        info="RDBT default: euler_ancestral. Unknown values fall back with a notice.",
                    )
                    sched = gr.Dropdown(
                        choices=list(config.SCHEDULER_CHOICES),
                        value=config.DEFAULT_SCHEDULER,
                        label="scheduler",
                        info="RDBT default: simple.",
                    )
                    denoise = gr.Slider(
                        config.MIN_DENOISE,
                        config.MAX_DENOISE,
                        value=config.DEFAULT_DENOISE,
                        step=config.DENOISE_STEP,
                        label="denoise",
                    )
                gallery = gr.Gallery(
                    label="output",
                    columns=2,
                    object_fit="contain",
                    height=600,
                )
                status = gr.Markdown(
                    "**While the Space starts:** the RDBT checkpoint downloads (and optional hub assets from config; no GPU needed). "
                    "**On first Generate:** ZeroGPU assigns a GPU, then the Diffusers pipeline loads and runs your job — "
                    "that first click can take extra time to build and load the model."
                )
                go = gr.Button("Generate", variant="primary")
        go.click(
            fn=_generate_task,
            inputs=[prompt, neg, w, h, st, cfg, batch, sampler, sched, denoise],
            outputs=[gallery, status],
        )
    return demo


# Hugging Face imports this module as `import app` (not __main__), so the block
# below must NOT be the only place that calls `queue()`. Otherwise the Space
# serves `demo` without a queue and shutdown can get noisy async teardown.
demo = build_ui()
demo.queue()

# Prepare Comfy + weights at container import time so the first HTTP request does
# not pay download cost; the Space "Starting…" state covers this (skip in tests).
if not config.skip_startup_bootstrap():
    from src.diffusers_backend import run_at_container_startup

    run_at_container_startup()

# Gradio 6: `theme` is a launch() argument. ZeroGPU calls `demo.launch()` without
# running our `if __name__` block, so wrap to always default the Soft theme.
_base_launch = demo.launch

@functools.wraps(_base_launch)  # type: ignore[misc]
def _launch_with_default_theme(*args: Any, **kwargs: Any):
    kwargs.setdefault("theme", gr.themes.Soft())
    return _base_launch(*args, **kwargs)


demo.launch = _launch_with_default_theme  # type: ignore[method-assign, assignment]

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False,
    )