"""detectFaces, trackFaces, faceLandmarks.

Detection is OpenCV YuNet (MIT). Tracking is a ByteTrack-style IoU tracker
implemented here: high-confidence detections associate first, low-confidence
ones only extend existing tracks, and a track ends after 500 ms without a match.
Identity embeddings (for local person clustering) come from an opt-in backend
and are null otherwise.
"""

from __future__ import annotations

import math
import threading
from dataclasses import dataclass, field
from typing import Any

import cv2
import numpy as np

from .. import media_io
from ..schemas import DetectFacesRequest, TimesRequest, TrackFacesRequest
from .common import Ctx, clamp01, norm_box, norm_point, open_input, unique_sorted

DETECT_MAX_SIDE = 1280
DETECT_THRESHOLD = 0.6
TRACK_LOW_THRESHOLD = 0.3
TRACK_MATCH_IOU = 0.3
TRACK_MAX_GAP_MS = 500


@dataclass
class Detection:
    """Pixel-space face: box (x, y, w, h), score, five landmarks (x, y)."""

    box: tuple[float, float, float, float]
    score: float
    landmarks: list[tuple[float, float]]


class YuNetDetector:
    def __init__(self, model_path: str) -> None:
        self.model_path = model_path
        self._local = threading.local()

    def _detector(self, width: int, height: int) -> Any:
        detector = getattr(self._local, "detector", None)
        if detector is None:
            detector = cv2.FaceDetectorYN.create(
                self.model_path, "", (width, height), TRACK_LOW_THRESHOLD, 0.3, 5000
            )
            self._local.detector = detector
        detector.setInputSize((width, height))
        return detector

    def detect(self, frame: np.ndarray) -> list[Detection]:
        height, width = frame.shape[:2]
        _, faces = self._detector(width, height).detect(frame)
        result: list[Detection] = []
        for row in faces if faces is not None else []:
            points = [(float(row[4 + 2 * i]), float(row[5 + 2 * i])) for i in range(5)]
            result.append(
                Detection(
                    (float(row[0]), float(row[1]), float(row[2]), float(row[3])), float(row[14]), points
                )
            )
        return result


def load_face_detector(files: dict[str, str], _settings: Any) -> YuNetDetector:
    return YuNetDetector(files["face_detection_yunet_2023mar.onnx"])


class SFaceEmbedder:
    def __init__(self, model_path: str) -> None:
        self.model_path = model_path
        self._local = threading.local()

    def embed(self, frame: np.ndarray, detection: Detection) -> np.ndarray:
        recognizer = getattr(self._local, "recognizer", None)
        if recognizer is None:
            recognizer = cv2.FaceRecognizerSF.create(self.model_path, "")
            self._local.recognizer = recognizer
        row = np.array(
            [*detection.box, *[c for p in detection.landmarks for c in p], detection.score],
            dtype=np.float32,
        )
        aligned = recognizer.alignCrop(frame, row)
        return recognizer.feature(aligned).reshape(-1).astype(np.float32)


def load_identity_embedder(files: dict[str, str], _settings: Any) -> SFaceEmbedder:
    return SFaceEmbedder(files["face_recognition_sface_2021dec.onnx"])


def _observation(det: Detection, width: int, height: int) -> dict[str, Any]:
    return {
        "box": norm_box(*det.box, width, height),
        "confidence": clamp01(det.score),
        "landmarks": [norm_point(x, y, width, height) for x, y in det.landmarks],
    }


def detect_run(ctx: Ctx, body: DetectFacesRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    detector = ctx.registry.runtime("detectFaces")
    frames = (
        media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE)
        if body.timesMs is not None
        else media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE)
    )
    out = []
    for t_ms, frame in frames:
        height, width = frame.shape[:2]
        faces = [
            _observation(d, width, height) for d in detector.detect(frame) if d.score >= DETECT_THRESHOLD
        ]
        out.append({"tMs": t_ms, "faces": faces})
    return {"frames": out}


def landmarks_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    detector = ctx.registry.runtime("faceLandmarks")
    out = []
    for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE):
        height, width = frame.shape[:2]
        faces = [
            {
                "box": norm_box(*d.box, width, height),
                "landmarks": [norm_point(x, y, width, height) for x, y in d.landmarks],
            }
            for d in detector.detect(frame)
            if d.score >= DETECT_THRESHOLD
        ]
        out.append({"tMs": t_ms, "faces": faces})
    return {"frames": out}


def iou(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
    ax1, ay1, ax2, ay2 = a[0], a[1], a[0] + a[2], a[1] + a[3]
    bx1, by1, bx2, by2 = b[0], b[1], b[0] + b[2], b[1] + b[3]
    iw = max(0.0, min(ax2, bx2) - max(ax1, bx1))
    ih = max(0.0, min(ay2, by2) - max(ay1, by1))
    inter = iw * ih
    union = a[2] * a[3] + b[2] * b[3] - inter
    return inter / union if union > 0 else 0.0


@dataclass
class Track:
    track_id: str
    box: tuple[float, float, float, float]
    velocity: tuple[float, float] = (0.0, 0.0)
    last_ms: int = 0
    keyframes: list[dict[str, Any]] = field(default_factory=list)
    # Best (score, frame crop source, detection) for identity embeddings.
    best: list[tuple[float, np.ndarray, Detection]] = field(default_factory=list)

    def predicted(self, t_ms: int) -> tuple[float, float, float, float]:
        dt = t_ms - self.last_ms
        return (
            self.box[0] + self.velocity[0] * dt,
            self.box[1] + self.velocity[1] * dt,
            self.box[2],
            self.box[3],
        )


class IouTracker:
    """Two-stage greedy IoU association (ByteTrack's idea without the Kalman filter)."""

    def __init__(self, max_gap_ms: int = TRACK_MAX_GAP_MS) -> None:
        self.max_gap_ms = max_gap_ms
        self.active: list[Track] = []
        self.finished: list[Track] = []
        self._next = 1

    def _associate(
        self, tracks: list[Track], dets: list[Detection], t_ms: int
    ) -> tuple[list[tuple[Track, Detection]], list[Track], list[Detection]]:
        pairs = sorted(
            (
                (iou(tr.predicted(t_ms), d.box), ti, di)
                for ti, tr in enumerate(tracks)
                for di, d in enumerate(dets)
            ),
            reverse=True,
        )
        used_t: set[int] = set()
        used_d: set[int] = set()
        matches: list[tuple[Track, Detection]] = []
        for score, ti, di in pairs:
            if score < TRACK_MATCH_IOU:
                break
            if ti in used_t or di in used_d:
                continue
            used_t.add(ti)
            used_d.add(di)
            matches.append((tracks[ti], dets[di]))
        rest_t = [t for i, t in enumerate(tracks) if i not in used_t]
        rest_d = [d for i, d in enumerate(dets) if i not in used_d]
        return matches, rest_t, rest_d

    def update(
        self, t_ms: int, dets: list[Detection], frame: np.ndarray | None, keep_crops: bool
    ) -> list[tuple[Track, Detection]]:
        expired = [t for t in self.active if t_ms - t.last_ms > self.max_gap_ms]
        self.finished.extend(expired)
        self.active = [t for t in self.active if t not in expired]
        high = [d for d in dets if d.score >= DETECT_THRESHOLD]
        low = [d for d in dets if TRACK_LOW_THRESHOLD <= d.score < DETECT_THRESHOLD]
        matches, rest_tracks, rest_high = self._associate(self.active, high, t_ms)
        low_matches, _, _ = self._associate(rest_tracks, low, t_ms)
        matches += low_matches
        for track, det in matches:
            dt = max(t_ms - track.last_ms, 1)
            track.velocity = ((det.box[0] - track.box[0]) / dt, (det.box[1] - track.box[1]) / dt)
            track.box = det.box
            track.last_ms = t_ms
        for det in rest_high:
            track = Track(f"t{self._next}", det.box, last_ms=t_ms)
            self._next += 1
            self.active.append(track)
            matches.append((track, det))
        if keep_crops and frame is not None:
            for track, det in matches:
                if det.score >= DETECT_THRESHOLD:
                    track.best.append((det.score, frame, det))
                    track.best.sort(key=lambda item: -item[0])
                    del track.best[3:]
        return matches

    def close(self) -> list[Track]:
        tracks = self.finished + self.active
        self.active, self.finished = [], []
        return sorted(tracks, key=lambda t: int(t.track_id[1:]))


def track_run(ctx: Ctx, body: TrackFacesRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    detector = ctx.registry.runtime("trackFaces")
    embedder = (
        ctx.registry.optional_runtime("faceIdentity")
        if ctx.registry.optional_enabled("faceIdentity")
        else None
    )
    tracker = IouTracker()
    for t_ms, frame in media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE):
        height, width = frame.shape[:2]
        matches = tracker.update(t_ms, detector.detect(frame), frame, embedder is not None)
        for track, det in matches:
            track.keyframes.append(
                {"tMs": t_ms, "box": norm_box(*det.box, width, height), "confidence": clamp01(det.score)}
            )
    tracks = []
    for track in tracker.close():
        embedding = None
        if embedder is not None and track.best:
            vectors = np.stack([embedder.embed(frame, det) for _, frame, det in track.best])
            mean = vectors.mean(axis=0)
            norm = float(np.linalg.norm(mean))
            if norm > 0 and math.isfinite(norm):
                embedding = [float(v) for v in mean / norm]
        tracks.append({"trackId": track.track_id, "keyframes": track.keyframes, "embedding": embedding})
    return {"sampleFps": body.sampleFps, "tracks": tracks}