"""
Fast face swap via InsightFace inswapper_128 — deterministic, no prompts, no diffusion.

This is the same pipeline used by roop / FaceSwapKahraman: detect the source face
once, then for every frame detect faces and paste the swapped face back. Runs on
CPU (onnxruntime); GFPGAN enhancement is applied by the caller if requested.
"""

from __future__ import annotations

from typing import Callable

import numpy as np
from huggingface_hub import hf_hub_download

_swapper = None
_analyzer = None

INSWAPPER_REPO = "ezioruan/inswapper_128.onnx"
INSWAPPER_FILE = "inswapper_128.onnx"


def _make_models(providers, ctx_id):
    import insightface
    from insightface.app import FaceAnalysis
    # Fast swap runs InsightFace on CPU on ZeroGPU (onnxruntime-gpu doesn't win over the
    # CPU onnxruntime insightface pulls in). To stay inside the GPU time window we cut the
    # per-frame cost hard: load ONLY detection + recognition (inswapper needs the target
    # face's keypoints from detection + the source embedding from recognition — the
    # genderage and 2D/3D landmark models aren't used), and detect at 320px not 640px.
    # Together that's a ~5-8x speedup vs the old 5-model / 640px setup.
    analyzer = FaceAnalysis(
        name="buffalo_l",
        allowed_modules=["detection", "recognition"],
        providers=providers,
    )
    analyzer.prepare(ctx_id=ctx_id, det_size=(320, 320))
    model_path = hf_hub_download(INSWAPPER_REPO, INSWAPPER_FILE)
    swapper = insightface.model_zoo.get_model(model_path, providers=providers)
    return analyzer, swapper


def _get_models():
    """Prefer the GPU (CUDAExecutionProvider) — CPU inswapper is ~10-50x slower and
    blows the ZeroGPU time window on anything but the shortest clip. Falls back to CPU
    if onnxruntime-gpu / CUDA isn't actually usable, so there's no regression."""
    global _swapper, _analyzer
    if _swapper is not None and _analyzer is not None:
        return _analyzer, _swapper

    cuda = False
    try:
        import onnxruntime as ort
        cuda = "CUDAExecutionProvider" in ort.get_available_providers()
    except Exception:
        cuda = False

    if cuda:
        try:
            _analyzer, _swapper = _make_models(
                ["CUDAExecutionProvider", "CPUExecutionProvider"], 0
            )
            print("[fastswap] InsightFace running on CUDA")
        except Exception as e:  # CUDA listed but not actually usable → fall back
            print(f"[fastswap] CUDA init failed ({e}); falling back to CPU")
            _analyzer = _swapper = None

    if _swapper is None or _analyzer is None:
        _analyzer, _swapper = _make_models(["CPUExecutionProvider"], -1)
        print("[fastswap] InsightFace running on CPU")

    return _analyzer, _swapper


def fast_swap_video(
    frames: np.ndarray,
    face_image,
    progress_cb: Callable[[float, str], None] | None = None,
) -> tuple[np.ndarray, str]:
    """
    Swap the reference face onto every frame.

    Args:
        frames:     uint8 [N, H, W, 3] RGB guide frames
        face_image: PIL image of the reference face

    Returns:
        (uint8 [N, H, W, 3] RGB swapped frames, status message)
    """
    analyzer, swapper = _get_models()

    src_bgr = np.array(face_image.convert("RGB"))[:, :, ::-1].copy()
    src_faces = analyzer.get(src_bgr)
    if not src_faces:
        return frames, "No face detected in the reference image — returning guide video unchanged."
    source_face = max(src_faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))

    out = np.empty_like(frames)
    swapped_count = 0
    N = len(frames)
    for i in range(N):
        frame_bgr = frames[i][:, :, ::-1].copy()
        faces = analyzer.get(frame_bgr)
        if faces:
            for tgt in faces:
                frame_bgr = swapper.get(frame_bgr, tgt, source_face, paste_back=True)
            swapped_count += 1
        out[i] = frame_bgr[:, :, ::-1]
        if progress_cb and (i % 8 == 0 or i == N - 1):
            progress_cb(i / max(N - 1, 1), f"Swapping faces… {i + 1}/{N}")

    if swapped_count == 0:
        return frames, "No faces detected in the guide video — returning it unchanged."
    return out, f"Swapped {swapped_count}/{N} frames."