"""Provider zero: animate the performer's OWN photo with LTX-Video running on
this Space's ZeroGPU — real GPU work, so it belongs under @spaces.GPU
(unlike the HTTP provider chain, which ZeroGPU aborts for idling).

Quota reality (HF docs, checked 2026-11-05): a free account gets 5 GPU-minutes
per DAY, charged to whoever runs the function. One call here (load + ~5s clip)
is budgeted at up to 300s of wall time but should use ~30-90 GPU-seconds on
the Blackwell slice. That buys a few HERO clips a day — chorus moments — not
a whole video. Design integration accordingly: LTX for the first N shots
only (default 3), chain/photo-zoom for the rest.

Stage 2 (separate spike): lip-sync (SadTalker-class talking head) so the
performer's mouth/head ride the actual vocals.
"""
import os
from pathlib import Path

import spaces

MODEL_ID = os.environ.get("LTX_MODEL", "Lightricks/LTX-Video")
_pipe = None


# Only the diffusers-layout folders LTXImageToVideoPipeline loads. The repo
# also carries standalone ComfyUI checkpoints at its root (tens of GB) — a
# bare snapshot_download pulls ALL of it and can fill the Space's disk,
# killing the container mid-call (suspected cause of the 2026-11-06 404s).
_ALLOW = ["model_index.json", "scheduler/*", "tokenizer/*", "text_encoder/*",
          "transformer/*", "vae/*"]


def ensure_weights():
    """Download the weights on CPU so GPU calls never burn their window (or
    the daily quota) on downloads. Stages are printed so the container log
    shows exactly how far a call gets."""
    from huggingface_hub import snapshot_download
    print("[ltx] ensure_weights: downloading pipeline files…", flush=True)
    path = snapshot_download(MODEL_ID, allow_patterns=_ALLOW)
    print(f"[ltx] ensure_weights: weights ready at {path}", flush=True)


def _load_pipe():
    global _pipe
    if _pipe is None:
        import torch
        from diffusers import LTXImageToVideoPipeline
        _pipe = LTXImageToVideoPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
        _pipe.to("cuda")
    return _pipe


# ZeroGPU laws (learned live 2026-11-06): spaces multiplies the requested
# duration by 1.5 when reserving (300→"450s", 180→"270s", 120→180s — all
# confirmed in error messages), and the desk checks the reservation against
# the CALLER's remaining daily quota. The anonymous pool had 136s left, so
# duration=85 (reserves ~128s) squeezes under it; authenticating with the
# owner's HF token raises the pool to the account's own 5 GPU-min/day.
# Generation is trimmed to fit the smaller window: 16 steps.
@spaces.GPU(duration=85)
def ltx_clip(image_path, prompt, seconds, out_path):
    """Animate `image_path` into a ~`seconds`-long clip at `out_path`.
    Returns the output path. Raises on any failure (caller falls back)."""
    from diffusers.utils import export_to_video, load_image

    print("[ltx] ltx_clip: loading pipeline…", flush=True)
    pipe = _load_pipe()
    print("[ltx] ltx_clip: pipeline on GPU, generating…", flush=True)
    image = load_image(str(image_path))
    frames_n = int(min(max(float(seconds), 2.0), 6.0) * 24)
    frames_n = (frames_n // 8) * 8 + 1  # LTX wants 8n+1 frame counts
    frames = pipe(
        image=image,
        prompt=str(prompt)[:400],
        negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted",
        width=704,
        height=480,
        num_frames=frames_n,
        num_inference_steps=16,
    ).frames[0]
    out_path = Path(out_path)
    export_to_video(frames, str(out_path), fps=24)
    print(f"[ltx] ltx_clip: wrote {out_path}", flush=True)
    return out_path