"""
BFS — Best Face Swap Video  ·  Hugging Face Space
"""

from __future__ import annotations

import os
import tempfile

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

from composer import compose_frames, crop_reserved_region
from video_utils import (
    compute_target_size,
    extract_audio,
    frames_for_duration,
    load_video_frames,
    resize_frames,
    save_video,
)

# ---------------------------------------------------------------------------
# GPU decorator — no-op locally, activates the ZeroGPU grant on HF Spaces
# ---------------------------------------------------------------------------
try:
    import spaces
    GPU = spaces.GPU
except ImportError:
    def GPU(fn=None, **kwargs):  # type: ignore
        return fn if fn is not None else lambda f: f

# ---------------------------------------------------------------------------
# Global model state (loaded once per worker)
# ---------------------------------------------------------------------------
_pipeline_state: dict | None = None

REGION_SIZE = 256
DEFAULT_FPS = 24.0
DEFAULT_DURATION = 5.0
DEFAULT_RESOLUTION = 768


def make_temp_file(suffix: str) -> str:
    f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    path = f.name
    f.close()
    return path


# ---------------------------------------------------------------------------
# Core processing function
# ---------------------------------------------------------------------------

@GPU(duration=300)
def generate(
    guide_video_path: str | None,
    face_image: Image.Image | None,
    prompt: str,
    duration: float,
    fps: float,
    lora_strength: float,
    seed: int,
    hf_token: str = "",
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str | None, str]:
    """
    Full head-swap pipeline:
      1. Load + resize guide video frames
      2. Compose chroma face strip (ReservedRegionFrameComposer)
      3. Run LTX-2.3 diffusion
      4. Crop face strip from output
      5. Mux original audio back in

    Returns (output_video_path, status_message).
    """
    global _pipeline_state

    # ---- validate inputs early ----
    if not guide_video_path:
        return None, "Please up