#!/usr/bin/env python3
"""Lokaler Face-Swap mit InsightFace + inswapper_128.

Verwendung:
  python3 swap.py <target_image> <source_face_image> <output_image>

- target_image:        Bild, dessen Gesicht ersetzt werden soll (Body/Scene)
- source_face_image:   Bild, aus dem das Gesicht entnommen wird
- output_image:        Pfad fuer das Ergebnis (PNG/JPG)

Optionale Env:
  FACE_INDEX_TARGET   Index des Gesichts im Ziel (default 0 = groesstes)
  FACE_INDEX_SOURCE   Index im Source (default 0)
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis


MODEL_PATH = Path(__file__).resolve().parent.parent / "models" / "inswapper_128.onnx"


def load_face_app() -> FaceAnalysis:
    app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
    app.prepare(ctx_id=0, det_size=(640, 640))
    return app


def load_swapper():
    if not MODEL_PATH.exists():
        raise FileNotFoundError(f"Inswapper model fehlt: {MODEL_PATH}")
    return insightface.model_zoo.get_model(str(MODEL_PATH), providers=["CPUExecutionProvider"])


def pick_face(faces, index: int):
    if not faces:
        return None
    faces_sorted = sorted(
        faces,
        key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]),
        reverse=True,
    )
    if index >= len(faces_sorted):
        return None
    return faces_sorted[index]


def main(argv: list[str]) -> int:
    if len(argv) < 4:
        print(__doc__)
        return 2

    target_path, source_path, out_path = argv[1], argv[2], argv[3]
    idx_t = int(os.environ.get("FACE_INDEX_TARGET", "0"))
    idx_s = int(os.environ.get("FACE_INDEX_SOURCE", "0"))

    target_img = cv2.imread(target_path)
    source_img = cv2.imread(source_path)
    if target_img is None:
        raise SystemExit(f"Kann Ziel nicht laden: {target_path}")
    if source_img is None:
        raise SystemExit(f"Kann Source nicht laden: {source_path}")

    print(f"[info] target={target_img.shape}, source={source_img.shape}")

    app = load_face_app()
    swapper = load_swapper()

    target_faces = app.get(target_img)
    source_faces = app.get(source_img)
    print(f"[info] gefundene Gesichter: target={len(target_faces)} source={len(source_faces)}")

    target_face = pick_face(target_faces, idx_t)
    source_face = pick_face(source_faces, idx_s)
    if target_face is None:
        raise SystemExit("Kein Gesicht im Zielbild gefunden")
    if source_face is None:
        raise SystemExit("Kein Gesicht im Source-Bild gefunden")

    result = swapper.get(target_img, target_face, source_face, paste_back=True)

    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    ok = cv2.imwrite(str(out), result)
    if not ok:
        raise SystemExit(f"Konnte Ergebnis nicht schreiben: {out}")
    print(f"[ok] geschrieben: {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))
