from collections import deque
from dataclasses import dataclass

from edge.models import Person, BBox, Keypoint


@dataclass
class TrackFrame:
    t: float
    bbox: BBox
    keypoints: list  # list[Keypoint]


@dataclass
class Track:
    track_id: int
    frames: deque

    @property
    def latest(self) -> TrackFrame:
        return self.frames[-1]


def _iou(a: BBox, b: BBox) -> float:
    ix1, iy1 = max(a.x1, b.x1), max(a.y1, b.y1)
    ix2, iy2 = min(a.x2, b.x2), min(a.y2, b.y2)
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    union = a.width * a.height + b.width * b.height - inter
    return inter / union if union else 0.0


class IouTracker:
    """경량 IoU 그리디 매칭 추적기. 실장비는 ByteTrack 어댑터로 교체 가능(동일 TrackerPort)."""

    def __init__(self, iou_thr: float = 0.3, max_age: float = 1.0, history: int = 90):
        self.iou_thr = iou_thr
        self.max_age = max_age
        self.history = history
        self._tracks: dict[int, Track] = {}
        self._last_seen: dict[int, float] = {}
        self._lost: list[Track] = []
        self._next_id = 1

    def update(self, persons: list[Person], t: float) -> list[Track]:
        self._lost = []
        unmatched = list(persons)
        # 기존 트랙과 그리디 IoU 매칭
        for tid, track in list(self._tracks.items()):
            best, best_iou = None, self.iou_thr
            for p in unmatched:
                iou = _iou(track.latest.bbox, p.bbox)
                if iou >= best_iou:
                    best, best_iou = p, iou
            if best is not None:
                unmatched.remove(best)
                track.frames.append(TrackFrame(t, best.bbox, best.keypoints))
                self._last_seen[tid] = t
        # 신규 트랙
        for p in unmatched:
            tid = self._next_id
            self._next_id += 1
            self._tracks[tid] = Track(tid, deque([TrackFrame(t, p.bbox, p.keypoints)],
                                                 maxlen=self.history))
            self._last_seen[tid] = t
        # 소실 판정
        for tid in list(self._tracks.keys()):
            if t - self._last_seen[tid] > self.max_age:
                self._lost.append(self._tracks.pop(tid))
                self._last_seen.pop(tid, None)
        return list(self._tracks.values())

    def lost(self) -> list[Track]:
        return list(self._lost)