"""FastH3-Omni · reference-to-video (+audio) on ZeroGPU.

Ordered image, video and audio references plus a prompt become a short video
with a matching 32 kHz soundtrack in eight fused forwards of the FastVideo
FastH3 OmniRef student. A plain prompt is wrapped into the model's six-section
H3 format; "Upsample prompt" rewrites it with the MiniMax-H3 prompt rewriter
instead. See fasth3_engine.py for the ZeroGPU + FastVideo integration notes.
"""
import contextvars
import os
import tempfile
from functools import cache
from pathlib import Path

import gradio as gr
from PIL import Image

try:
    import spaces
    IS_ZERO = True
except ImportError:  # local / CI
    IS_ZERO = False

    class _NoSpaces:
        @staticmethod
        def GPU(*a, **k):
            if a and callable(a[0]):
                return a[0]
            return lambda f: f

    spaces = _NoSpaces()  # type: ignore

from fasth3_engine import FastH3OmniRefEngine, MODEL_REPO, prefetch

REWRITER_SPACE = os.getenv("FASTH3_REWRITER", "hugging-apps/minimax-h3-prompt-rewriter")

CANVASES = {
    "1344x768 · 16:9 recommended": (768, 1344),
    "832x480 · 16:9 fast": (480, 832),
    "1280x704 · 16:9": (704, 1280),
    "960x544 · 16:9": (544, 960),
    "768x1344 · 9:16 recommended": (1344, 768),
    "480x832 · 9:16 fast": (832, 480),
    "768x768 · 1:1": (768, 768),
    "1024x1024 · 1:1": (1024, 1024),
    "1024x768 · 4:3": (768, 1024),
    "768x1024 · 3:4": (1024, 768),
    "1536x672 · 21:9": (672, 1536),
}
DEFAULT_CANVAS = "1344x768 · 16:9 recommended"
PORTRAIT_CANVAS = "768x1344 · 9:16 recommended"
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
MIN_FRAMES, MAX_FRAMES = 124, 362
MIN_DURATION, MAX_UI_DURATION = 5, 14
MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
REWRITER_IMAGE_SLOTS = 4
DEFAULT_PIXEL_FRAMES = 1344 * 768 * 124
GPU_SECONDS_WARM = int(os.getenv("FASTH3_GPU_SECONDS", "170"))
GPU_SECONDS_COLD = int(os.getenv("FASTH3_GPU_SECONDS_COLD", "360"))
GPU_SECONDS_PER_EXTRA_IMAGE, GPU_SECONDS_PER_VIDEO = 110, 160

APP_DIR = Path(__file__).resolve().parent
EXAMPLES_DIR = APP_DIR / "examples"
LIMITS = {"image": 9, "video": 3, "audio": 3}
REWRITER_LIMITS = {"image": 4, "video": 1, "audio": 1}

EXAMPLES = [
    {
        "title": "Beach promenade at golden hour",
        "image": "woman.jpg",
        "desc": "a young woman with a dark curly afro and gold hoop earrings, wearing a white t-shirt",
        "scene": "walks along a sunlit beach promenade at golden hour, smiling and glancing toward the camera, "
        "palm trees and the ocean behind her.",
        "sound": "ocean waves, a gentle breeze and distant seagulls",
        "music": "none",
    },
    {
        "title": "Dog bounding through snow",
        "image": "dog.jpg",
        "desc": "a pale cream, fluffy long-haired dog with floppy ears and a dark nose",
        "scene": "bounds across a snowy field in bright winter sunlight, kicking up powder snow, then shakes itself off.",
        "sound": "crunching snow, happy panting and a cold wind",
        "music": "none",
    },
    {
        "title": "Birthday cake with candles",
        "image": "cake.jpg",
        "desc": "a whole chocolate layer cake with chocolate frosting, pastel sprinkles and chocolate shards on top",
        "scene": "sits at the center of a festive birthday table in a warm, dim room as lit candles flicker and "
        "colorful confetti drifts down around it.",
        "sound": "soft party chatter and a faint crackle",
        "music": "cheerful, upbeat birthday music",
    },
]


def composer_examples():
    rows = []
    for example in EXAMPLES:
        path = str(EXAMPLES_DIR / example["image"])
        rows.append({
            "title": example["title"],
            "refs": [{"kind": "image", "path": path, "url": f"/gradio_api/file={path}", "name": example["image"],
                      "desc": example["desc"]}],
            "builder": {"scene": example["scene"], "sound": example["sound"], "music": example["music"],
                        "style": "realistic photographic"},
        })
    return rows


def snap_frames(seconds: float) -> int:
    frames = max(1, round(float(seconds) * FPS))
    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
        frames += 1
    return frames


def reference_labels(references):
    names = {"image": "Picture", "video": "Video", "audio": "Audio"}
    counts, labels = {}, []
    for kind, _ in references:
        counts[kind] = counts.get(kind, 0) + 1
        labels.append((kind, f"<{names[kind]} {counts[kind]}>"))
    return labels


def to_h3_prompt(prompt: str, references) -> str:
    """Wrap a plain prompt into the six-section H3 format the checkpoint was trained on."""
    prompt = (prompt or "").strip()
    if "subject_definitions:" in prompt:
        return prompt
    labels = reference_labels(references)
    visual = " and ".join(label for kind, label in labels if kind != "audio") or "the references"
    audio = [label for kind, label in labels if kind == "audio"]
    sound = (f"the voice and sound of {' and '.join(audio)}, in a natural acoustic space matching the scene"
             if audio else "natural ambient sound that matches the scene")
    return (
        f"subject_definitions:\n<Subject 1> is the subject shown in {visual}.\n\n"
        f"summary:\n{prompt}\n\n"
        f"retention_analysis:\n<Subject 1>: fully_preserved - the reference subject's identity, "
        f"appearance and details are kept.\n\n"
        f"detailed_description:\nThe target video is in a realistic photographic style.\n"
        f"[Shot 1] {prompt}\n\n"
        f"overall_soundscape:\n{sound}.\n\n"
        f"non_diegetic_music:\nnone."
    )


# --- engine: built once at startup, on ZeroGPU's emulated CUDA ------------
# The weights land on "cuda" in the main process; ZeroGPU packs them and maps
# them onto the real card at the start of every GPU call, so no call rebuilds.
from fasth3_engine import fetch_kernel_cache, zerogpu_main_process_shims

prefetch()
KERNELS_CACHED = fetch_kernel_cache()
if IS_ZERO:
    zerogpu_main_process_shims()
ENGINE = FastH3OmniRefEngine()
ENGINE.initialize()
_COMPLETED_GPU_CALLS = 0


def probe(path: str):
    import av

    def seconds(stream, container):
        if stream.duration is not None and stream.time_base is not None:
            return float(stream.duration * stream.time_base)
        return None if container.duration is None else container.duration / av.time_base

    with av.open(path) as container:
        video = seconds(container.streams.video[0], container) if container.streams.video else None
        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
    return video, audio


def _allowed(path: str) -> bool:
    from gradio.utils import get_upload_folder

    real = os.path.realpath(path)
    roots = [os.path.realpath(get_upload_folder()), os.path.realpath(EXAMPLES_DIR)]
    return any(real == root or real.startswith(root + os.sep) for root in roots) and os.path.isfile(real)


def references_from(composer: dict):
    refs = (composer or {}).get("refs") or []
    if any(ref.get("uploading") for ref in refs):
        raise gr.Error("A reference is still uploading. Give it a second and try again.")
    for ref in refs:
        if ref.get("error"):
            raise gr.Error(ref["error"])
    references = []
    for ref in refs:
        kind, path = ref.get("kind"), ref.get("path")
        if kind not in LIMITS or not path or not _allowed(path):
            raise gr.Error("One of the references could not be read. Remove it and add it again.")
        references.append((kind, _as_rgb_png(path, len(references)) if kind == "image" else path))
    for kind, limit in LIMITS.items():
        if sum(k == kind for k, _ in references) > limit:
            raise gr.Error(f"At most {limit} {kind} references.")
    return references


def _as_rgb_png(path: str, index: int) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=f"_ref{index + 1}.png", delete=False)
    Image.open(path).convert("RGB").save(tmp.name)
    return tmp.name


def soundtrack_seconds(references):
    carried = []
    for kind, path in references:
        if kind == "image":
            continue
        _, audio_seconds = probe(path)
        if audio_seconds is not None:
            carried.append(audio_seconds)
    return carried


def check(prompt: str, references) -> None:
    if not references:
        raise gr.Error("Add at least one image or video reference.")
    if not any(kind != "audio" for kind, _ in references):
        raise gr.Error("Audio needs an image or a video alongside it.")
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the video: fill in the builder or write a prompt.")
    if "<removed>" in prompt:
        raise gr.Error("The prompt still refers to a removed reference (<removed>). Edit it out first.")
    for kind, path in references:
        if kind != "video":
            continue
        video_seconds, _ = probe(path)
        if video_seconds is None:
            raise gr.Error("A reference video has no video stream. Add it as audio instead.")
        if not MIN_REFERENCE_VIDEO <= video_seconds <= MAX_REFERENCE_VIDEO:
            raise gr.Error(
                f"A reference video is {video_seconds:.1f} s. Use clips between "
                f"{MIN_REFERENCE_VIDEO:g} and {MAX_REFERENCE_VIDEO:g} seconds."
            )


@cache
def rewriter():
    from gradio_client import Client

    return Client(REWRITER_SPACE, token=False)


def aspect_label(height: int, width: int) -> str:
    ratio = width / height
    options = {"21:9": 21 / 9, "16:9": 16 / 9, "4:3": 4 / 3, "1:1": 1.0, "3:4": 3 / 4, "9:16": 9 / 16}
    return min(options, key=lambda key: abs(options[key] - ratio))


def resolve_frames(references, match, duration) -> int:
    carried = soundtrack_seconds(references)
    if match and len(carried) == 1 and MIN_FRAMES <= snap_frames(carried[0]) <= MAX_FRAMES:
        return snap_frames(carried[0])
    return min(MAX_FRAMES, max(MIN_FRAMES, snap_frames(duration)))


def upsample(composer, canvas, match, duration):
    from gradio_client import handle_file

    composer = dict(composer or {})
    try:
        references = references_from(composer)
        if not any(kind != "audio" for kind, _ in references):
            raise gr.Error("Add an image or a video reference first.")
        for kind, limit in REWRITER_LIMITS.items():
            if sum(k == kind for k, _ in references) > limit:
                raise gr.Error(
                    f"The rewriter takes at most {REWRITER_LIMITS['image']} images, 1 video and 1 audio clip."
                )
        prompt = (composer.get("upsample_source") or "").strip()
        if not prompt:
            raise gr.Error("Describe the video first, then upsample it.")
        missing = [label for _, label in reference_labels(references) if label not in prompt]
        if missing:
            prompt = f"Keep the subject shown in {' and '.join(missing)}. {prompt}"
        height, width = CANVASES.get(canvas, CANVASES[DEFAULT_CANVAS])
        num_frames = resolve_frames(references, match, duration)
        images = [path for kind, path in references if kind == "image"]
        images += [None] * (REWRITER_LIMITS["image"] - len(images))
        video = next((path for kind, path in references if kind == "video"), None)
        audio = next((path for kind, path in references if kind == "audio"), None)
        rewritten, status = rewriter().predict(
            task="Ref2AV",
            prompt=prompt,
            duration=round(num_frames / FPS, 2),
            resolution=aspect_label(height, width),
            image_1=handle_file(images[0]) if images[0] else None,
            image_2=handle_file(images[1]) if images[1] else None,
            image_3=handle_file(images[2]) if images[2] else None,
            image_4=handle_file(images[3]) if images[3] else None,
            ref_video=handle_file(video) if video else None,
            ref_audio=handle_file(audio) if audio else None,
            api_name="/rewrite",
        )
        rewritten = (rewritten or "").strip()
        if not rewritten:
            raise gr.Error(f"The rewriter returned nothing: {status}")
    except gr.Error as error:
        gr.Warning(str(error.message))
        rewritten = ""
    except Exception as error:
        gr.Warning(f"The prompt rewriter failed: {type(error).__name__}: {error}")
        rewritten = ""
    composer["upsampled"] = rewritten
    return composer


def _gpu_duration(prompt, references, height, width, num_frames, seed, *args, **kwargs):
    cold = _COMPLETED_GPU_CALLS == 0 and not KERNELS_CACHED
    base = GPU_SECONDS_COLD if cold else GPU_SECONDS_WARM
    seconds = base * max(1.0, (height * width * num_frames) / DEFAULT_PIXEL_FRAMES)
    images = sum(kind == "image" for kind, _ in references)
    seconds += GPU_SECONDS_PER_EXTRA_IMAGE * max(0, images - 1)
    seconds += GPU_SECONDS_PER_VIDEO * sum(kind == "video" for kind, _ in references)
    return int(min(900, seconds))


@spaces.GPU(duration=_gpu_duration, size="xlarge")
def _generate(prompt, references, height, width, num_frames, seed, progress=gr.Progress()):
    progress(0.01, desc="Starting on the GPU")
    context = contextvars.copy_context()

    def report(fraction, desc):
        context.copy().run(progress, fraction, desc=desc)

    return ENGINE.generate(prompt, references, height, width, num_frames, seed, on_progress=report)


def generate(composer, canvas=DEFAULT_CANVAS, match=True, duration=5, seed=1234, progress=gr.Progress()):
    references = references_from(composer)
    prompt = ((composer or {}).get("prompt") or "").strip()
    check(prompt, references)
    height, width = CANVASES.get(canvas, CANVASES[DEFAULT_CANVAS])
    num_frames = resolve_frames(references, match, duration)
    seed = int(seed) if seed is not None else 0
    global _COMPLETED_GPU_CALLS
    progress(0, desc="Waiting for a GPU")
    path = _generate(to_h3_prompt(prompt, references), references, height, width, num_frames, seed)
    _COMPLETED_GPU_CALLS += 1
    return path


# --- UI ---------------------------------------------------------------------
HEADER = """
<div class="fh-header">
  <h1>FastH3-<span class="fh-accent">Omni</span></h1>
  <p>FastH3-Omni is FastVideo's 8-step reference-to-video student that turns reference images, videos and
  audio plus a prompt into a short clip with a matching soundtrack. Drop in references, then describe the video
  with the builder or write the prompt yourself, tagging each reference by name.</p>
  <div class="fh-pills">
    <a href="https://huggingface.co/FastVideo/FastVideo-FastH3-Omni-8-Step-V1" target="_blank">Model</a>
    <a href="https://huggingface.co/collections/FastVideo/fasth3-omni-6ad501acfccdebd3c3b7445b" target="_blank">Collection</a>
    <a href="https://github.com/hao-ai-lab/FastVideo" target="_blank">FastVideo</a>
    <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank">MiniMax-H3</a>
  </div>
</div>
"""

FASTVIDEO_BLUE = gr.themes.Color(
    c50="#eef3ff", c100="#dbe5ff", c200="#bccfff", c300="#8eaeff", c400="#5f8cff",
    c500="#346cff", c600="#1f55ec", c700="#1843c4", c800="#17389c", c900="#16317c", c950="#0e1d4a",
    name="fastvideo",
)

CSS = """
.main.fillable { max-width: 1250px !important; }
.fh-header { padding: 6px 2px 12px; }
.fh-header h1 { font-size: 1.6rem; margin: 0 0 6px; font-weight: 700; }
.fh-header h1 .fh-accent { color: #346cff; }
.fh-header p { color: var(--body-text-color-subdued); margin: 0 0 12px; max-width: 76ch; }
.fh-pills { display: flex; gap: 8px; flex-wrap: wrap; }
.fh-pills a { font-size: .8rem; padding: 3px 11px; border-radius: 999px;
  background: var(--button-secondary-background-fill); color: var(--body-text-color);
  text-decoration: none; border: 1px solid var(--border-color-primary); }
.fh-pills a:hover { border-color: #346cff; color: #346cff; }
#gen-btn { background: #346cff; color: #fff; border: none; }
#gen-btn:hover { background: #1f55ec; }
"""

THEME = gr.themes.Base(
    primary_hue=FASTVIDEO_BLUE,
    secondary_hue=FASTVIDEO_BLUE,
    neutral_hue="zinc",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
).set(
    body_background_fill_dark="#000000",
    block_background_fill_dark="#0b0b0f",
    background_fill_secondary_dark="#0b0b0f",
    color_accent="#346cff",
    slider_color="#346cff",
    checkbox_background_color_selected="#346cff",
    checkbox_background_color_selected_dark="#346cff",
)

COMPOSER_HTML = (APP_DIR / "composer.html").read_text()
COMPOSER_CSS = (APP_DIR / "composer.css").read_text()
COMPOSER_JS = (APP_DIR / "composer.js").read_text()

with gr.Blocks(title="FastH3-Omni") as demo:
    gr.HTML(HEADER)

    with gr.Row():
        with gr.Column(scale=1):
            composer = gr.HTML(
                value={"examples": composer_examples()},
                html_template=COMPOSER_HTML,
                css_template=COMPOSER_CSS,
                js_on_load=COMPOSER_JS,
                apply_default_css=False,
            )
            run = gr.Button("Generate", variant="primary", size="lg", elem_id="gen-btn")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                match = gr.Checkbox(
                    label="Match the reference soundtrack",
                    info="When exactly one reference carries a soundtrack, the clip takes its length.",
                    value=True,
                )
                duration = gr.Slider(
                    label="Duration (s)", minimum=MIN_DURATION, maximum=MAX_UI_DURATION, step=1, value=5
                )
                seed = gr.Number(label="Seed", value=1234, precision=0)
                gr.Markdown("8 fused steps and guidance 1.0 are fixed by the model's contract.")

        with gr.Column(scale=1):
            result = gr.Video(label="Video + soundtrack", autoplay=True)

    composer.edit(upsample, [composer, canvas, match, duration], [composer], api_name=False, show_progress="hidden")
    run.click(generate, [composer, canvas, match, duration, seed], [result], api_name="generate")


if __name__ == "__main__":
    demo.queue(max_size=12).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        theme=THEME,
        css=CSS,
        ssr_mode=False,
        allowed_paths=[str(EXAMPLES_DIR)],
    )