from __future__ import annotations

import asyncio.base_events
import gc
import os
import random
import threading
import time
import traceback
import warnings
from dataclasses import dataclass
from pathlib import Path


_ORIGINAL_EVENT_LOOP_DEL = asyncio.base_events.BaseEventLoop.__del__


def _quiet_invalid_fd_event_loop_del(self, _warn=warnings.warn):
    try:
        _ORIGINAL_EVENT_LOOP_DEL(self, _warn)
    except ValueError as exc:
        if "Invalid file descriptor" not in str(exc):
            raise


asyncio.base_events.BaseEventLoop.__del__ = _quiet_invalid_fd_event_loop_del
warnings.filterwarnings(
    "ignore",
    message=r".*HTTP_422_UNPROCESSABLE_ENTITY.*",
    category=Warning,
    module=r"gradio\.routes",
)

import gradio as gr
import spaces
import torch

from sefi import SEFIInferencePipeline


os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

CACHE_DIR = os.getenv(
    "SEFI_CACHE_DIR",
    "/data/sefi-cache" if os.path.isdir("/data") else "/tmp/sefi-cache",
)
APP_DIR = Path(__file__).resolve().parent
EXAMPLE_ASSETS_DIR = APP_DIR / "assets" / "examples"
gr.set_static_paths([EXAMPLE_ASSETS_DIR])


@dataclass(frozen=True)
class ModelPreset:
    label: str
    repo_id: str
    family: str
    steps: int
    guidance: float


MODEL_PRESETS: dict[str, ModelPreset] = {
    "1b-base": ModelPreset(
        label="SeFi-Image 1B Base",
        repo_id="SeFi-Image/SeFi-Image-1B-Base",
        family="base",
        steps=50,
        guidance=4.0,
    ),
    "2b-base": ModelPreset(
        label="SeFi-Image 2B Base",
        repo_id="SeFi-Image/SeFi-Image-2B-Base",
        family="base",
        steps=50,
        guidance=4.0,
    ),
    "5b-base": ModelPreset(
        label="SeFi-Image 5B Base",
        repo_id="SeFi-Image/SeFi-Image-5B-Base",
        family="base",
        steps=50,
        guidance=4.0,
    ),
    "5b-rl": ModelPreset(
        label="SeFi-Image 5B RL",
        repo_id="SeFi-Image/SeFi-Image-5B-RL",
        family="rl",
        steps=50,
        guidance=4.0,
    ),
    "1b-turbo": ModelPreset(
        label="SeFi-Image 1B Turbo",
        repo_id="SeFi-Image/SeFi-Image-1B-turbo",
        family="turbo",
        steps=4,
        guidance=1.0,
    ),
    "2b-turbo": ModelPreset(
        label="SeFi-Image 2B Turbo",
        repo_id="SeFi-Image/SeFi-Image-2B-turbo",
        family="turbo",
        steps=4,
        guidance=1.0,
    ),
    "5b-turbo": ModelPreset(
        label="SeFi-Image 5B Turbo",
        repo_id="SeFi-Image/SeFi-Image-5B-turbo",
        family="turbo",
        steps=4,
        guidance=1.0,
    ),
}

DEFAULT_MODEL = "1b-turbo"
APP_TITLE = "SeFi-Image"
MODEL_PUBLISHER_URL = "https://huggingface.co/SeFi-Image"
PROJECT_PAGE_URL = "https://jmliu206.github.io/sefi-web/"
IMAGE_SIZE = 1024
TURBO_STEPS = {4, 8, 10}
GPU_DURATION_BY_MODEL = {
    "1b-turbo": 25,
    "2b-turbo": 35,
    "5b-turbo": 65,
    "1b-base": 40,
    "2b-base": 60,
    "5b-base": 110,
    "5b-rl": 110,
}
EXTRA_STEP_SECONDS_BY_MODEL = {
    "1b-turbo": 2,
    "2b-turbo": 3,
    "5b-turbo": 5,
    "1b-base": 1,
    "2b-base": 1,
    "5b-base": 2,
    "5b-rl": 2,
}
EXAMPLES = [
    {
        "label": "Anime cinematic portrait",
        "model_key": "5b-rl",
        "prompt": (
            "Anime-realistic cinematic portrait of a young adult woman in a "
            "rain-lit Tokyo alley, expressive eyes, subtle natural skin texture, "
            "wind-touched dark hair, layered streetwear jacket, neon reflections, "
            "shallow depth of field, 85mm lens look, elegant color grading, "
            "highly detailed face, painterly anime realism, dramatic rim light, "
            "clean composition, masterpiece quality."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7401,
        "image": str(EXAMPLE_ASSETS_DIR / "anime_cinematic_portrait_1024.png"),
    },
    {
        "label": "Realistic editorial portrait",
        "model_key": "5b-rl",
        "prompt": (
            "Realistic editorial portrait of a confident young adult man in a "
            "black turtleneck and tailored coat, soft window light, textured "
            "gray studio backdrop, cinematic shadows, detailed eyes, natural "
            "skin, medium-format fashion photography, restrained luxury mood, "
            "precise facial anatomy, sharp focus, subtle film grain, balanced "
            "composition."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7402,
        "image": str(EXAMPLE_ASSETS_DIR / "realistic_editorial_portrait_1024.png"),
    },
    {
        "label": "Mythic alpine landscape",
        "model_key": "5b-rl",
        "prompt": (
            "Incredible mythic alpine landscape at sunrise, enormous "
            "snow-covered mountains above a crystal lake, floating mist, "
            "wildflowers in the foreground, tiny stone observatory on a ridge, "
            "golden light breaking through storm clouds, epic depth, sweeping "
            "cinematic composition, ultra detailed environment, atmospheric "
            "perspective, realistic fantasy concept art, awe inspiring scale."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7403,
        "image": str(EXAMPLE_ASSETS_DIR / "mythic_alpine_landscape_1024.png"),
    },
    {
        "label": "Cyrillic museum poster",
        "model_key": "5b-rl",
        "prompt": (
            'A complex museum poster for a northern light exhibition, strong '
            'editorial typography, large readable Cyrillic headline "СЕВЕРНЫЙ '
            'СВЕТ", smaller Cyrillic subheading "выставка света и льда", '
            "asymmetric Swiss grid composition, layered translucent paper, icy "
            "blue and black ink, precise margins, high-end graphic design, "
            "photographed as a printed poster on a gallery wall."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7301,
        "image": str(EXAMPLE_ASSETS_DIR / "cyrillic_museum_poster_1024.png"),
    },
    {
        "label": "Korean night jazz poster",
        "model_key": "5b-rl",
        "prompt": (
            'A sophisticated Korean night jazz festival poster, large Hangul '
            'title "서울의 밤", smaller Hangul text "재즈 페스티벌", vertical '
            "composition with a saxophone silhouette, neon reflections on rain, "
            "black paper, magenta and cyan spot colors, tight typographic "
            "hierarchy, centered moon circle, elegant modern Seoul design."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7302,
        "image": str(EXAMPLE_ASSETS_DIR / "korean_night_jazz_poster_1024.png"),
    },
    {
        "label": "Chinese tea packaging",
        "model_key": "5b-rl",
        "prompt": (
            'Premium Chinese tea packaging poster, large brush-style Chinese '
            'characters "春风茶馆", small vertical Chinese seal text, ceramic '
            "tea cup and folded paper wrapper, balanced negative space, deep "
            "jade green and warm ivory, gold foil accents, product photography "
            "mixed with graphic layout, calm luxury composition."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7303,
        "image": str(EXAMPLE_ASSETS_DIR / "chinese_tea_packaging_1024.png"),
    },
    {
        "label": "Multiscript design festival",
        "model_key": "5b-rl",
        "prompt": (
            'International design festival poster with three writing systems in '
            'one composition: Cyrillic "ТИХИЙ ГОРОД", Korean "고요한 도시", '
            'Chinese "静城". Use the texts as bold typographic blocks, modular '
            "grid, architectural isometric city fragments, layered risograph "
            "texture, red black and pale gray palette, disciplined composition, "
            "poster photographed flat on a studio table."
        ),
        "steps": 50,
        "guidance": 4.0,
        "seed": 7304,
        "image": str(EXAMPLE_ASSETS_DIR / "multiscript_design_festival_1024.png"),
    },
]

_MODEL_LOCK = threading.Lock()
_LOADED_MODEL_KEY: str | None = None
_LOADED_PIPE: SEFIInferencePipeline | None = None


def _model_choices() -> list[tuple[str, str]]:
    return [(preset.label, key) for key, preset in MODEL_PRESETS.items()]


def _example_samples() -> list[list[str]]:
    samples = []
    for example in EXAMPLES:
        preset = MODEL_PRESETS[str(example["model_key"])]
        settings = (
            f"{IMAGE_SIZE}x{IMAGE_SIZE}, {example['steps']} steps, "
            f"guidance {example['guidance']}, seed {example['seed']}"
        )
        samples.append(
            [
                str(example["image"]),
                str(example["label"]),
                preset.label,
                str(example["prompt"]),
                settings,
            ]
        )
    return samples


def _torch_cleanup() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def _clear_loaded_model() -> None:
    global _LOADED_MODEL_KEY, _LOADED_PIPE
    _LOADED_PIPE = None
    _LOADED_MODEL_KEY = None
    _torch_cleanup()


def _load_pipe(model_key: str) -> SEFIInferencePipeline:
    global _LOADED_MODEL_KEY, _LOADED_PIPE

    preset = MODEL_PRESETS[model_key]
    with _MODEL_LOCK:
        if _LOADED_PIPE is not None and _LOADED_MODEL_KEY == model_key:
            return _LOADED_PIPE

        _clear_loaded_model()
        pipe = SEFIInferencePipeline.from_pretrained(
            preset.repo_id,
            cache_dir=CACHE_DIR,
            device="cuda",
            dtype="bf16",
        )
        _LOADED_MODEL_KEY = model_key
        _LOADED_PIPE = pipe
        return pipe


def _friendly_error(exc: BaseException, repo_id: str | None = None) -> str:
    text = str(exc)
    lowered = text.lower()
    gated = (
        "requires approval" in lowered
        or "gated" in lowered
        or "401" in lowered
        or "403" in lowered
    )
    if gated:
        repo_hint = f" for `{repo_id}`" if repo_id else ""
        return (
            f"Model access is not approved{repo_hint}. Open the model page while "
            "logged in as the Space owner, accept the SeFi non-commercial gate, "
            "and retry. The Space already has `HF_TOKEN` configured as a secret."
        )

    return f"{type(exc).__name__}: {text}"


def _format_seconds(seconds: float) -> str:
    seconds = max(0, int(round(seconds)))
    minutes, secs = divmod(seconds, 60)
    if minutes:
        return f"{minutes}m {secs:02d}s"
    return f"{secs}s"


def _format_timing(seconds: float) -> str:
    seconds = max(0.0, float(seconds))
    if seconds < 10:
        return f"{seconds:.1f}s"
    return _format_seconds(seconds)


def _format_generation_summary(
    *,
    preset: ModelPreset,
    width: int,
    height: int,
    steps: int,
    guidance_scale: float,
    seed: int,
    total_seconds: float,
    model_load_seconds: float,
    model_was_loaded: bool,
    setup_seconds: float,
    denoise_seconds: float,
    decode_seconds: float,
) -> str:
    if denoise_seconds > 0 and steps > 0:
        speed = f"{steps / denoise_seconds:.2f} steps/s ({denoise_seconds / steps:.2f}s/step)"
    else:
        speed = "n/a"
    load_note = "already in memory" if model_was_loaded else "download/load/switch"

    return (
        "### Generation summary\n\n"
        f"`{preset.repo_id}` · {width}x{height} · {steps} steps · "
        f"guidance {guidance_scale} · seed {seed}\n\n"
        "| Phase | Time |\n"
        "| --- | ---: |\n"
        f"| Total backend time | {_format_timing(total_seconds)} |\n"
        f"| Model {load_note} | {_format_timing(model_load_seconds)} |\n"
        f"| Prompt + latent setup | {_format_timing(setup_seconds)} |\n"
        f"| Raw denoising | {_format_timing(denoise_seconds)} |\n"
        f"| Decode + output | {_format_timing(decode_seconds)} |\n"
        f"| Raw denoising speed | {speed} |\n"
    )


def model_defaults(model_key: str):
    preset = MODEL_PRESETS[model_key]
    return (
        gr.update(value=preset.steps),
        gr.update(value=preset.guidance),
        (
            f"Selected `{preset.repo_id}`. Defaults: "
            f"{preset.steps} steps, guidance {preset.guidance}."
        ),
    )


def load_example(index: int):
    example = EXAMPLES[int(index)]
    preset = MODEL_PRESETS[str(example["model_key"])]
    status = (
        f'Loaded example "{example["label"]}" generated with `{preset.repo_id}` '
        f"at {IMAGE_SIZE}x{IMAGE_SIZE}, {example['steps']} steps, "
        f'guidance {example["guidance"]}, seed {example["seed"]}.'
    )
    return (
        example["model_key"],
        example["prompt"],
        example["steps"],
        example["guidance"],
        example["seed"],
        True,
        example["image"],
        status,
    )


def estimate_duration(
    model_key: str,
    prompt: str,
    steps: int,
    guidance_scale: float,
    seed: int,
    randomize_seed: bool,
    *_args,
    **_kwargs,
) -> int:
    del prompt, guidance_scale, seed, randomize_seed
    preset = MODEL_PRESETS.get(model_key)
    duration = GPU_DURATION_BY_MODEL.get(model_key, 60)
    if preset is not None:
        extra_steps = max(0, int(steps) - preset.steps)
        duration += extra_steps * EXTRA_STEP_SECONDS_BY_MODEL.get(model_key, 2)
    return duration


@spaces.GPU(duration=estimate_duration)
def generate(
    model_key: str,
    prompt: str,
    steps: int,
    guidance_scale: float,
    seed: int,
    randomize_seed: bool,
    progress=gr.Progress(track_tqdm=False),
):
    request_started_at = time.monotonic()
    prompt = prompt.strip()
    if not prompt:
        return None, "Enter a prompt.", seed

    preset = MODEL_PRESETS[model_key]
    steps = int(steps)
    guidance_scale = float(guidance_scale)
    width = IMAGE_SIZE
    height = IMAGE_SIZE

    if preset.family == "turbo":
        if steps not in TURBO_STEPS:
            return (
                None,
                "Turbo checkpoints currently support 4, 8, or 10 denoising steps.",
                seed,
            )
        if guidance_scale != 1.0:
            return None, "Turbo checkpoints should use guidance 1.0.", seed

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)

    try:
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True

        progress(0, desc=f"Loading {preset.label}")
        model_was_loaded = _LOADED_PIPE is not None and _LOADED_MODEL_KEY == model_key
        load_started_at = time.monotonic()
        pipe = _load_pipe(model_key)
        load_finished_at = time.monotonic()
        pipe_started_at = load_finished_at
        denoise_started_at: float | None = None
        last_denoise_step_at: float | None = None
        progress(0, desc="Preparing prompt and latents")

        def report_step(step: int, total: int) -> None:
            nonlocal denoise_started_at, last_denoise_step_at
            total = max(1, int(total))
            step = min(max(0, int(step)), total)
            now = time.monotonic()
            if step == 0:
                denoise_started_at = now
            if denoise_started_at is None:
                denoise_started_at = now
            if step > 0:
                last_denoise_step_at = now
            elapsed = now - denoise_started_at
            remaining = 0.0
            if step > 0:
                remaining = (elapsed / step) * (total - step)
            desc = (
                f"Denoising {step}/{total} steps | "
                f"elapsed {_format_seconds(elapsed)} | "
                f"ETA {_format_seconds(remaining)}"
            )
            progress(step / total, desc=desc)

        images = pipe(
            prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            seed=int(seed),
            progress_callback=report_step,
        )
        pipe_finished_at = time.monotonic()
        progress(1, desc="Finalizing image")
    except Exception as exc:
        traceback.print_exc()
        return None, _friendly_error(exc, preset.repo_id), seed

    if not images:
        return None, "Generation finished without an image.", seed

    request_finished_at = time.monotonic()
    if denoise_started_at is None:
        denoise_started_at = pipe_started_at
    if last_denoise_step_at is None:
        last_denoise_step_at = denoise_started_at

    return (
        images[0],
        _format_generation_summary(
            preset=preset,
            width=width,
            height=height,
            steps=steps,
            guidance_scale=guidance_scale,
            seed=seed,
            total_seconds=request_finished_at - request_started_at,
            model_load_seconds=load_finished_at - load_started_at,
            model_was_loaded=model_was_loaded,
            setup_seconds=denoise_started_at - pipe_started_at,
            denoise_seconds=last_denoise_step_at - denoise_started_at,
            decode_seconds=pipe_finished_at - last_denoise_step_at,
        ),
        seed,
    )


APP_CSS = """
#examples_table {
    width: 100%;
}

#examples_table .table-wrap,
#examples_table .table-wrap.fixed-height {
    height: auto !important;
    max-height: none !important;
}

#examples_table table {
    width: 100%;
}

#examples_table th,
#examples_table td {
    vertical-align: middle;
}

#examples_table img {
    width: 64px !important;
    height: 64px !important;
    object-fit: cover;
    border-radius: 6px;
}

#examples_table td {
    white-space: normal;
    line-height: 1.35;
}

#generation_status table {
    width: 100%;
    margin-top: 0.5rem;
}

#generation_status th,
#generation_status td {
    padding: 6px 8px;
}

#generation_status th:last-child,
#generation_status td:last-child {
    text-align: right;
    white-space: nowrap;
}
"""


with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(
        f"""
# {APP_TITLE}

SeFi-Image text-to-image checkpoints in Base, Turbo, and 5B RL variants. Base
models use the 50-step guidance-4.0 setting, Turbo models use the 4-step
guidance-1.0 setting, and 5B RL is the reinforced-learning checkpoint.

[Models on Hugging Face]({MODEL_PUBLISHER_URL}) | [Project page]({PROJECT_PAGE_URL})
"""
    )

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            model = gr.Dropdown(
                label="Model",
                choices=_model_choices(),
                value=DEFAULT_MODEL,
                interactive=True,
            )
            prompt = gr.Textbox(
                label="Prompt",
                value="A blue ceramic mug on a white desk.",
                lines=4,
                max_lines=8,
            )
            with gr.Row():
                steps = gr.Slider(
                    minimum=1,
                    maximum=60,
                    step=1,
                    value=MODEL_PRESETS[DEFAULT_MODEL].steps,
                    label="Steps",
                )
                guidance = gr.Slider(
                    minimum=1.0,
                    maximum=8.0,
                    step=0.1,
                    value=MODEL_PRESETS[DEFAULT_MODEL].guidance,
                    label="Guidance",
                )
            with gr.Row():
                seed = gr.Number(
                    label="Seed",
                    value=42,
                    precision=0,
                    minimum=0,
                    maximum=2**31 - 1,
                )
                randomize_seed = gr.Checkbox(label="Randomize", value=True)
            with gr.Row():
                run = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1, min_width=360):
            image = gr.Image(label="Image", type="pil", format="png")
            status = gr.Markdown(
                (
                    f"Selected `{MODEL_PRESETS[DEFAULT_MODEL].repo_id}`. Defaults: "
                    f"{MODEL_PRESETS[DEFAULT_MODEL].steps} steps, "
                    f"guidance {MODEL_PRESETS[DEFAULT_MODEL].guidance}."
                ),
                elem_id="generation_status",
            )

    examples = gr.Dataset(
        samples=_example_samples(),
        components=[
            gr.Image(label="Image", type="filepath", height=64, width=64, render=False),
            gr.Textbox(label="Example", render=False),
            gr.Textbox(label="Model", render=False),
            gr.Textbox(label="Prompt", render=False),
            gr.Textbox(label="Settings", render=False),
        ],
        headers=["Image", "Example", "Model", "Prompt", "Settings"],
        type="index",
        layout="table",
        label="Examples",
        samples_per_page=10,
        elem_id="examples_table",
    )

    model.change(model_defaults, inputs=model, outputs=[steps, guidance, status])
    examples.click(
        load_example,
        inputs=examples,
        outputs=[
            model,
            prompt,
            steps,
            guidance,
            seed,
            randomize_seed,
            image,
            status,
        ],
        api_name="load_example",
        api_visibility="undocumented",
    )
    run.click(
        generate,
        inputs=[
            model,
            prompt,
            steps,
            guidance,
            seed,
            randomize_seed,
        ],
        outputs=[image, status, seed],
        api_name="generate",
        concurrency_limit=1,
    )

demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(
        allowed_paths=[str(EXAMPLE_ASSETS_DIR)],
        css=APP_CSS,
        ssr_mode=False,
    )