"""Sprite Forge — game sprite sheets from one character design, via MiniMax-H3.

An implementation of the workflow Victor Mustar described: design one character on a flat magenta backdrop, generate
one short H3 clip per move, cut each clip down to its pose extremes, key out the magenta, and pack the result.

Three things are specific to doing it here rather than by hand:

* **Every clip loops.** Each move is generated with the character passed as *both* the first and the last keyframe,
  so H3's `fl2va` branch has to leave the rest pose and arrive back at it. A sprite that does not return to its start
  pose visibly pops every time the animation repeats.
* **The moves are generated concurrently.** They are independent conditioning-identical requests, so the wall clock
  is one clip rather than the sum of them.
* **This Space owns no GPU.** It runs on `cpu-basic` and calls other Spaces for both models. See `backends.py`.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import time
import traceback

import gradio as gr
import numpy as np
from PIL import Image

import backends
import moves as move_defs
import player
import spritekit as sk

OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "sprite-forge")
os.makedirs(OUTPUT_DIR, exist_ok=True)

DEFAULT_SUBJECT = "a scrappy alley cat in a bomber jacket"
EXAMPLES = [
    "a scrappy alley cat in a bomber jacket",
    "a round little wizard with an oversized pointed hat",
    "a rusty barrel-chested robot with one glowing eye",
    "a punk frog in a leather vest",
    "a knight in dented armor holding a short sword",
]

PENDING = (
    "<div style='display:flex;align-items:center;justify-content:center;height:560px;opacity:.55;"
    "border:1px dashed var(--border-color-primary);border-radius:8px'>Designing your character…</div>"
)


IDLE_PANEL = (
    "<div style='display:flex;align-items:center;justify-content:center;height:560px;opacity:.55;"
    "border:1px dashed var(--border-color-primary);border-radius:8px'>"
    "Your playable character will appear here.</div>"
)


def building_panel(character_path: str | None, message: str) -> str:
    """The output panel while the clips are still generating: the character, and what is happening to it.

    The still is the first thing worth looking at and it exists a couple of minutes before the playable demo does,
    so it goes straight into the main panel rather than staying folded away in the downloads accordion.

    Embedded as a data URI, downscaled first. A file path here would have to go through `/gradio_api/file=`, which
    is the route that already 403'd on this Space once, and the panel is re-yielded on every step — a 1024px PNG
    would be a megabyte of base64 each time.
    """
    if not character_path:
        return PENDING
    try:
        thumbnail = Image.open(character_path).convert("RGBA")
        thumbnail.thumbnail((460, 460))
        source = sk.data_uri(thumbnail)
    except Exception:  # noqa: BLE001 — the panel is cosmetic; never fail a run over it
        return PENDING
    return (
        "<div style='display:flex;flex-direction:column;align-items:center;justify-content:center;gap:14px;"
        "height:560px;border:1px dashed var(--border-color-primary);border-radius:8px'>"
        "<style>@keyframes sf-pulse{0%,100%{opacity:.55}50%{opacity:1}}</style>"
        f"<img src='{source}' alt='character' style='max-height:400px;max-width:90%;border-radius:6px'>"
        f"<div style='animation:sf-pulse 1.6s ease-in-out infinite;font-size:.95em'>{message}</div>"
        "</div>"
    )


def _out(name: str) -> str:
    return os.path.join(OUTPUT_DIR, f"{int(time.time() * 1000)}-{name}")


def _adopt(path: str, name: str) -> str:
    """Copy a file fetched from a remote Space into this Space's own output directory.

    `gradio_client` drops downloads in gradio's temp cache, and `/gradio_api/file=` only serves whitelisted
    directories — handing a component a path under someone else's cache returns 403 to the browser. Copying makes
    the file ours to serve, and incidentally keeps it alive past gradio's cache cleanup.
    """
    destination = _out(name)
    shutil.copyfile(path, destination)
    return destination


# ---------------------------------------------------------------------------
# the one flow
# ---------------------------------------------------------------------------


def forge(
    subject,
    uploaded,
    style,
    backend_label,
    canvas_label,
    selected,
    duration,
    steps,
    seed,
    tolerance,
    softness,
    despill,
    despeckle,
    shrink,
    palette_size,
    flatten,
    playback_fps,
    sprite_height,
    request: gr.Request,
    progress=gr.Progress(),
):
    """Character -> one clip per move -> keyed extremes -> packed sheet -> playable demo, as one generator.

    Yields on every step so the demo panel holds Gradio's pending state from the moment the button is pressed rather
    than only lighting up at the end of a multi-minute run.
    """
    if not uploaded and not (subject or "").strip():
        raise gr.Error("Describe a character, or upload one.")
    if not selected:
        raise gr.Error("Pick at least one move under Advanced.")

    ip_token = backends.ip_token_of(request)
    presets = [move_defs.BY_LABEL[label] for label in selected]
    moves_label = f"{len(presets)} move" + ("s" if len(presets) != 1 else "")

    # ---- 1. the character design still ----
    if uploaded:
        character = uploaded
        message = f"Animating {moves_label}…"
        yield building_panel(character, message), character, [], None, [], None, "Using your uploaded character."
    else:
        yield PENDING, None, [], None, [], None, "Designing the character…"
        backend = backends.BY_LABEL[backend_label]
        _, width, height = backends.CANVASES[canvas_label]
        try:
            character = _adopt(
                backend.generate(
                    move_defs.character_prompt(subject, style), width, height, int(seed), ip_token
                ),
                "character.png",
            )
        except Exception as error:
            traceback.print_exc()
            raise gr.Error(f"{backend.space} failed: {type(error).__name__}: {error}") from error
        yield (
            building_panel(character, f"Animating {moves_label}…"),
            character, [], None, [], None, f"Character designed. Animating {moves_label}…",
        )

    # ---- 1b. flatten the still before it is animated ----
    # Snapping the still to a small palette *before* H3 sees it attacks the artifact at the source rather than
    # cleaning it up afterwards: the backdrop becomes one exact value, and the encoder gets large constant-color
    # regions instead of gradients to smear. The palette is then reused verbatim by the keyer.
    palette = None
    if int(palette_size) > 0:
        still = np.asarray(Image.open(character).convert("RGB"), np.float32) / 255.0
        still_key = sk.detect_key_color(still)
        palette = sk.extract_palette(sk.sample_interior(still, still_key, 8000), int(palette_size), still_key)
        print(f"[forge] palette: {len(palette)} colors from the character still", flush=True)
        if bool(flatten) and len(palette):
            flat = sk.quantize_to_palette(still, palette, still_key)
            character = _out("character-flat.png")
            sk.to_pil(np.concatenate([flat, np.ones((*flat.shape[:2], 1), np.float32)], axis=-1)).convert(
                "RGB"
            ).save(character)
            yield (
                building_panel(character, f"Flattened to {len(palette)} colors. Animating {moves_label}…"),
                character, [], None, [], None, f"Flattened to {len(palette)} colors.",
            )

    # ---- 2. one clip per move, concurrently ----
    # These requests are fully independent, so they are fired together: the wall clock becomes the slowest single
    # clip instead of the sum. Each gets its own seed so two moves never collapse onto the same motion.
    who = subject if not uploaded else "the character"

    def submit(preset, index, minimal=False):
        return backends.animate_job(
            prompt=preset.prompt(who, minimal=minimal),
            image_path=character,
            canvas_label=canvas_label,
            duration=float(duration),
            steps=int(steps),
            seed=int(seed) + index + 1,
            ip_token=ip_token,
        )

    def await_all(pending, label):
        """Poll a batch of in-flight jobs, driving the progress bar off the slowest one."""
        started = time.time()
        while not all(job.done() for _p, _i, job in pending):
            finished = sum(1 for _p, _i, job in pending if job.done())
            elapsed = time.time() - started
            # All the requests start together, so the batch takes as long as its slowest member and elapsed time is
            # a fair proxy for how far along that one is. The completed count is the ground truth and always wins;
            # the clock only fills the gap between completions, and is capped short of 1 so the bar never claims to
            # be finished before the work is.
            clock = min(elapsed / max(ANIMATE_SECONDS, 1.0), 0.95)
            fraction = max(finished / len(pending), clock)
            progress(0.1 + 0.65 * fraction, desc=f"{label} — {finished}/{len(pending)} done")
            time.sleep(0.4)
        return pending

    ANIMATE_SECONDS = 18 + 2.5 * int(steps) + 5.0 * float(duration)

    progress(0.1, desc=f"Animating {moves_label}")
    inflight = [(preset, i, submit(preset, i)) for i, preset in enumerate(presets)]
    await_all(inflight, f"Animating {moves_label}")

    clips, failures, refused = [], [], []
    for preset, index, job in inflight:
        try:
            video, report = backends.unwrap(job.result())
        except Exception as error:  # noqa: BLE001 — reported per move in the status line
            # The H3 Space runs a content classifier over the prompt of every keyframed request, and it is touchy in
            # ways the wording does not predict — "the foot sweeping up and out" is refused where "a front kick to
            # the right" is not. It runs *before* the GPU booking, so a refusal costs no quota and retrying is cheap.
            if "content filter" in str(error):
                refused.append((preset, index))
            else:
                failures.append(f"**{preset.label}**: `{type(error).__name__}: {error}`")
            continue
        clips.append({
            "key": preset.key,
            "label": preset.label,
            "video": _adopt(video, f"{preset.key}.mp4"),
            "frames": preset.frames,
            "cyclic": preset.cyclic,
        })

    if refused:
        # One retry wave, plainer phrasing and no character description.
        print(f"[forge] retrying {[p.key for p, _ in refused]} minimally", flush=True)
        retries = [(preset, i, submit(preset, i, minimal=True)) for preset, i in refused]
        await_all(retries, "Retrying refused moves")
        for preset, _index, job in retries:
            try:
                video, report = backends.unwrap(job.result())
            except Exception as error:  # noqa: BLE001
                failures.append(f"**{preset.label}**: `{type(error).__name__}: {error}`")
                continue
            clips.append({
                "key": preset.key,
                "label": preset.label,
                "video": _adopt(video, f"{preset.key}.mp4"),
                "frames": preset.frames,
                "cyclic": preset.cyclic,
            })
        # Keep sheet order matching the order the moves were requested in.
        order = {preset.key: i for i, preset in enumerate(presets)}
        clips.sort(key=lambda c: order[c["key"]])

    if not clips:
        raise gr.Error("Every move failed. " + (failures[0] if failures else ""))

    gallery = [(clip["video"], clip["label"]) for clip in clips]
    yield (
        building_panel(character, "Cutting to pose extremes…"),
        character, gallery, None, [], None, f"{len(clips)} clip" + ("s" if len(clips) != 1 else "") + " generated. Cutting to pose extremes…",
    )

    # ---- 3. key, cut to extremes, pack ----
    # `palette` was read in 1b, from the still and never from the clips. The still is a PNG: no chroma subsampling,
    # and a backdrop that is genuinely one flat color. H3's backdrop drifts slightly across a frame, and a darker
    # patch of magenta carries *less* chroma, so the interior test mistakes it for character — sampling the clips
    # put a near-backdrop color in the palette and turned whole backgrounds opaque.
    selected_frames, key_hex = [], None
    for i, clip in enumerate(clips):
        progress(0.75 + 0.15 * i / len(clips), desc=f"Keying {clip['label']}")
        frames, _fps = sk.read_video(clip["video"])
        keyed, key_rgb, reliable = sk.key_clip(
            frames,
            tolerance=float(tolerance),
            softness=float(softness),
            despill=float(despill),
            despeckle=int(despeckle),
            shrink=float(shrink),
            palette=palette,
        )
        key_hex = "#%02X%02X%02X" % tuple((key_rgb * 255).astype(int))

        chooser = sk.cycle_extremes if clip.get("cyclic") else sk.pose_extremes
        indices = chooser(keyed, int(clip["frames"]), candidates=reliable)
        selected_frames.append((clip, [keyed[j] for j in indices]))
        del keyed, frames

    # One crop box and therefore one scale, across every frame of every move.
    #
    # Cropping per move and scaling each to a fixed height is the obvious thing and it is wrong: a crouch has a
    # smaller bounding box, so normalising it to the same height renders it *larger* than the idle stance. A single
    # shared box also keeps the character in one coordinate system, so a jump genuinely rises inside its cell, a
    # crouch genuinely sinks, and a punch whose arm widens the silhouette does not shove the body sideways relative
    # to the other rows.
    progress(0.92, desc="Packing the sheet")
    box = sk.content_bbox([f for _clip, group in selected_frames for f in group])
    packed, previews = [], []
    for clip, group in selected_frames:
        sprites = sk.fit_height([sk.to_pil(f).crop(box) for f in group], int(sprite_height))
        packed.append(sk.Move(name=clip["key"], frames=sprites, fps=float(playback_fps)))
        preview = _out(f"{clip['key']}.webp")
        sk.save_animation(sprites, preview, fps=float(playback_fps))
        previews.append((preview, f"{clip['label']} · {len(sprites)} extremes"))

    sheet, atlas = sk.pack_sheet(packed)
    atlas["source"] = {"model": "MiniMaxAI/MiniMax-H3", "lora": "larryvrh/MiniMax-H3-Turbo-Lora", "keyed": key_hex}

    sheet_path = _out("spritesheet.png")
    sheet.save(sheet_path)
    atlas_path = sk.write_atlas(atlas, _out("atlas.json"))

    demo_html = player.build_demo(atlas, sk.data_uri(sheet), title="Sprite demo")
    demo_path = _out("demo.html")
    with open(demo_path, "w", encoding="utf-8") as handle:
        handle.write(demo_html)

    total = sum(len(m.frames) for m in packed)
    status = (
        f"**{len(packed)} moves · {total} frames** · cell {atlas['cell']['width']}×{atlas['cell']['height']} · "
        f"sheet {atlas['sheet']['width']}×{atlas['sheet']['height']} · keyed `{key_hex}` · "
        f"click the demo, then use the keys."
    )
    if failures:
        status += "\n\nFailed: " + "; ".join(failures)

    yield (
        player.as_iframe(demo_html),
        character,
        gallery,
        sheet_path,
        previews,
        [sheet_path, atlas_path, demo_path],
        status,
    )


# ---------------------------------------------------------------------------
# ui
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1300px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

# Gradio 6 moved `theme` and `css` off the Blocks constructor onto `launch()`; passing them here is only a warning,
# so the theme silently does not apply.
with gr.Blocks(title="Sprite Forge") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🕹️ Sprite Forge
            Describe a character, get a **playable** one. Powered by
            [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) with the 6-step
            [Turbo LoRA](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora), after
            [@victormustar's workflow](https://x.com/victormustar/status/2089310616854892818).

            One clip per move, each generated with the same first *and* last frame so it loops — then keyed, cut to
            its pose extremes and packed into a sprite sheet. Runs on your own ZeroGPU quota.
            """
        )

        with gr.Row(equal_height=False):
            # Controls left, the thing you actually came for on the right.
            with gr.Column(scale=2):
                subject = gr.Textbox(
                    label="Character",
                    value=DEFAULT_SUBJECT,
                    placeholder="a scrappy alley cat in a bomber jacket",
                    lines=2,
                )
                run = gr.Button("Generate", variant="primary", size="lg")
                gr.Examples(EXAMPLES, inputs=subject, label="Try a character")

                with gr.Accordion("Advanced", open=False):
                    uploaded = gr.Image(
                        label="Or start from your own character", type="filepath", height=200
                    )
                    selected_moves = gr.CheckboxGroup(
                        move_defs.LABELS, value=move_defs.DEFAULT_MOVES, label="Moves"
                    )
                    with gr.Row():
                        style = gr.Dropdown(
                            list(move_defs.STYLES), value=move_defs.DEFAULT_STYLE, label="Style"
                        )
                        image_backend = gr.Dropdown(
                            backends.BACKEND_LABELS, value=backends.BACKEND_LABELS[0], label="Image model"
                        )
                    with gr.Row():
                        canvas = gr.Dropdown(
                            list(backends.CANVASES), value=backends.DEFAULT_CANVAS, label="Canvas"
                        )
                        seed = gr.Number(value=1234, label="Seed", precision=0)
                    with gr.Row():
                        duration = gr.Slider(
                            2, 5, value=3, step=0.5, label="Clip seconds",
                            info="H3 does 5s per move; 3s keeps each move punchy.",
                        )
                        steps = gr.Slider(4, 8, value=6, step=1, label="H3 steps")
                    with gr.Row():
                        playback_fps = gr.Slider(4, 24, value=12, step=1, label="Playback fps")
                        sprite_height = gr.Slider(
                            128, 768, value=512, step=32, label="Max sprite height",
                            info="Only ever downscales. 512 keeps a 544px clip at native size.",
                        )
                    with gr.Row():
                        tolerance = gr.Slider(0.02, 0.30, value=0.10, step=0.01, label="Key tolerance")
                        softness = gr.Slider(0.01, 0.25, value=0.09, step=0.01, label="Edge softness")
                    with gr.Row():
                        despill = gr.Slider(0, 1, value=1.0, step=0.05, label="Despill")
                        despeckle = gr.Slider(0, 4, value=2, step=1, label="Despeckle")
                        palette_size = gr.Slider(
                            0, 32, value=16, step=2, label="Palette colors",
                            info="Flat-art keying: 0 falls back to plain chroma keying.",
                        )
                        shrink = gr.Slider(
                            0, 0.5, value=0.3, step=0.05, label="Matte shrink",
                            info="Chokes the matte inward against 4:2:0 chroma bleed.",
                        )
                    flatten = gr.Checkbox(
                        value=True, label="Flatten the character before animating",
                        info="Snaps the still to its palette first, so H3 animates genuinely flat art.",
                    )

            with gr.Column(scale=3):
                demo_frame = gr.HTML(IDLE_PANEL, label="Playable")
                status = gr.Markdown()

        with gr.Accordion("Sprite sheet, clips and downloads", open=False):
            with gr.Row():
                character_preview = gr.Image(label="Character still", height=250)
                clip_gallery = gr.Gallery(label="H3 clips", columns=3, height=250, object_fit="contain")
            sheet_image = gr.Image(label="Sprite sheet (transparent PNG)", height=250)
            sprite_previews = gr.Gallery(label="Looping moves", columns=5, height=210, object_fit="contain")
            downloads = gr.Files(label="Download — sheet, atlas.json, standalone demo.html")

    run.click(
        forge,
        [
            subject, uploaded, style, image_backend, canvas, selected_moves, duration, steps, seed,
            tolerance, softness, despill, despeckle, shrink, palette_size, flatten, playback_fps, sprite_height,
        ],
        [demo_frame, character_preview, clip_gallery, sheet_image, sprite_previews, downloads, status],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=4).launch(
        theme=gr.themes.Citrus(),
        css=CSS,
        ssr_mode=False,
        show_error=True,
        allowed_paths=[OUTPUT_DIR],
    )