"""
LTX-Video-0.9.8-13B-distilled inference pipeline for BFS head-swap.

Full pipeline loaded from a single version-matched repo:
  Lightricks/LTX-Video-0.9.8-13B-distilled
  (VAE, scheduler, text_encoder, tokenizer, and transformer are all 0.9.8)

LoRA:
  Alissonerdx/BFS-Best-Face-Swap-Video — ltx-2/head_swap_v2_multimodes.safetensors

Hardware note (ZeroGPU):
  This Space runs on ZeroGPU (H200 slice, ~70 GB VRAM). The pipeline MUST be
  loaded at module import time in app.py — ZeroGPU forks a fresh process for
  every @spaces.GPU call, so anything loaded lazily inside a GPU function is
  thrown away when the call ends and reloaded on every click.
  `pipe.to("cuda")` at startup is the official ZeroGPU pattern; the spaces
  package virtualizes CUDA until a GPU is actually attached.
"""

from __future__ import annotations

import gc
import os
from typing import Callable

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

PIPELINE_REPO = "Lightricks/LTX-Video-0.9.8-13B-distilled"
BFS_REPO = "Alissonerdx/BFS-Best-Face-Swap-Video"
BFS_FILE = "ltx-2/head_swap_v2_multimodes.safetensors"

NEGATIVE_PROMPT = (
    "pc game, console game, video game, cartoon, childish, ugly, "
    "artifacts, low resolution, blurry, jagged edges"
)


def _get_pipeline_cls():
    """Return the best available LTX I2V pipeline class."""
    try:
        from diffusers import LTXConditionPipeline
        return LTXConditionPipeline
    except ImportError:
        from diffusers import LTXImageToVideoPipeline
        return LTXImageToVideoPipeline


def load_pipeline(
    device: str = "cuda",
    token: str | None = None,
    progress_cb: Callable[[str], None] | None = None,
) -> dict:
    """
    Load the full 0.9.8-matched pipeline from the distilled repo.

    All components (VAE, scheduler, text_encoder, tokenizer, transformer)
    come from a single repo so they are guaranteed version-consistent.

    LoRA is loaded but NOT fused — lora_strength is applied per-request
    via set_adapters() in run_inference().
    """
    effective_token = token or os.environ.get("HF_TOKEN")

    PipelineCls = _get_pipeline_cls()

    if progress_cb:
        progress_cb(f"Loading full pipeline from {PIPELINE_REPO} ({PipelineCls.__name__})…")

    pipe = PipelineCls.from_pretrained(
        PIPELINE_REPO,
        torch_dtype=torch.bfloat16,
        token=effective_token,
    )

    pipe.to(device)

    if progress_cb:
        progress_cb("Loading BFS head-swap LoRA (ltx-2 variant, 48 layers)…")
    bfs_path = hf_hub_download(
        repo_id=BFS_REPO,
        filename=BFS_FILE,
        token=effective_token,
    )
    pipe.load_lora_weights(bfs_path, adapter_name="bfs")
    pipe.set_adapters(["bfs"], adapter_weights=[1.0])

    supports_video_condition = PipelineCls.__name__ == "LTXConditionPipeline"
    return {"pipe": pipe, "supports_video_condition": supports_video_condition}


def run_inference(
    state: dict,
    composed_frames: np.ndarray,
    prompt: str,
    fps: float = 24.0,
    lora_strength: float = 1.0,
    seed: int = 42,
    num_inference_steps: int = 8,
    guidance_scale: float = 1.0,
    condition_mode: str = "Guide video (V2V)",
    condition_strength: float = 0.7,
    denoise_strength: float = 1.0,
    region_size_px: int = 256,
    progress_cb: Callable[[str], None] | None = None,
) -> np.ndarray:
    """
    Run 13B-distilled LTX inference on the composed frames.

    BFS V3 is a persistent-template *video-to-video* workflow: the composed
    guide video (chroma strip on every frame) is passed as a video condition,
    and denoise_strength < 1.0 re-renders it with the swapped head while
    preserving the guide motion. "First frame only (I2V)" is kept as a
    fallback mode — motion then comes from the prompt alone.

    Args:
        state:              dict returned by load_pipeline()
        composed_frames:    uint8 [N, H, W, 3] with chroma strip composited in
        prompt:             text prompt (head_swap: format)
        fps:                target frame rate
        lora_strength:      BFS LoRA weight (0.0–2.0)
        seed:               RNG seed
        num_inference_steps: 4–8 typical for distilled model
        guidance_scale:     1.0 recommended for distilled (CFG disabled)
        condition_mode:     "Guide video (V2V)" or "First frame only (I2V)"
        condition_strength: conditioning strength for the guide video
        denoise_strength:   V2V only — how much to re-render (1.0 = from scratch)
        region_size_px:     strip width (informational — not used in pipe call)

    Returns:
        uint8 [N, H, W, 3] — generated frames (strip still present;
        call composer.crop_reserved_region() to remove it)
    """
    pipe = state["pipe"]
    use_video_condition = (
        state.get("supports_video_condition", False)
        and condition_mode.startswith("Guide video")
    )

    N, H, W, _ = composed_frames.shape

    pipe.set_adapters(["bfs"], adapter_weights=[lora_strength])

    generator = torch.Generator(device="cuda").manual_seed(seed)

    common = dict(
        prompt=prompt,
        negative_prompt=NEGATIVE_PROMPT,
        width=W,
        height=H,
        num_frames=N,
        frame_rate=int(fps),
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        decode_timestep=0.05,
        decode_noise_scale=0.025,
        output_type="np",
    )

    if progress_cb:
        progress_cb(f"Running diffusion ({'V2V' if use_video_condition else 'I2V'})…")

    with torch.inference_mode():
        if use_video_condition:
            # Full composed guide video as condition (frame count is 8k+1,
            # guaranteed by video_utils.frames_for_duration).
            video_pil = [Image.fromarray(f) for f in composed_frames]
            result = pipe(
                video=video_pil,
                frame_index=0,
                strength=condition_strength,
                denoise_strength=denoise_strength,
                **common,
            )
        else:
            result = pipe(
                image=Image.fromarray(composed_frames[0]),
                **common,
            )

    # output_type="np" → result.frames[0] is [N, H, W, C] float32 in [0, 1]
    frames_np = (result.frames[0] * 255).clip(0, 255).astype(np.uint8)

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return frames_np