"""目标跟踪服务 (MVP-2 / D19)

基于 ByteTrack 思想的简化跟踪器，纯 Python 实现，无需额外依赖。

核心算法 (ByteTrack 论文要点):

1. 按置信度将检测分为 high (>= high_thresh) / low (< high_thresh)
2. 先用 high 与现有 tracks 进行 IOU 匹配 (匈牙利算法 / 贪心)
3. 未匹配的 tracks 再与 low 检测匹配 (拯救低置信度真实目标)
4. 仍未匹配的 high 检测创建新 track
5. 未匹配的 tracks 进入 Lost 状态，超过 max_lost_frames 移除

简化点 (相对原版):

- 不使用 Kalman Filter，只用上一帧 bbox 与当前检测 IOU 匹配
- 匹配算法用贪心 (按 IOU 降序) 替代匈牙利
- 不支持外观特征 (ReID)

集成方式::

    tracker = ByteTracker()
    for frame_detections in stream:
        tracked = tracker.update(frame_detections)
        # tracked 中每个 UnifiedDetection 会被填充 track_id
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple

from models.event_schemas import BBox, UnifiedDetection

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Track 状态
# ---------------------------------------------------------------------------


class TrackState(str, Enum):
    """跟踪状态。"""

    NEW = "new"
    TRACKED = "tracked"
    LOST = "lost"
    REMOVED = "removed"


@dataclass
class Track:
    """单个目标跟踪。"""

    track_id: int
    bbox: BBox
    confidence: float
    class_name: str
    state: TrackState = TrackState.NEW
    age: int = 0  # 总存活帧数
    lost_frames: int = 0  # 连续丢失帧数
    last_update_frame: int = 0


# ---------------------------------------------------------------------------
# 工具函数: IOU
# ---------------------------------------------------------------------------


def compute_iou(b1: BBox, b2: BBox) -> float:
    """计算两个 BBox 的 IOU。"""

    inter_x1 = max(b1.x1, b2.x1)
    inter_y1 = max(b1.y1, b2.y1)
    inter_x2 = min(b1.x2, b2.x2)
    inter_y2 = min(b1.y2, b2.y2)

    inter_w = max(0, inter_x2 - inter_x1)
    inter_h = max(0, inter_y2 - inter_y1)
    inter = inter_w * inter_h

    if inter == 0:
        return 0.0

    area1 = b1.area
    area2 = b2.area
    union = area1 + area2 - inter
    return inter / union if union > 0 else 0.0


def greedy_match(
    cost_matrix: List[List[float]],
    threshold: float,
) -> List[Tuple[int, int]]:
    """贪心匹配: 按 IOU 降序选择不冲突的最佳匹配。

    Args:
        cost_matrix: cost_matrix[i][j] = IOU 越大越好
        threshold: 低于该阈值不匹配

    Returns:
        匹配对列表 [(row_idx, col_idx), ...]
    """

    matches: List[Tuple[int, int]] = []
    if not cost_matrix or not cost_matrix[0]:
        return matches

    candidates: List[Tuple[float, int, int]] = []
    for i, row in enumerate(cost_matrix):
        for j, score in enumerate(row):
            if score >= threshold:
                candidates.append((score, i, j))

    candidates.sort(reverse=True)
    used_rows: set = set()
    used_cols: set = set()
    for score, i, j in candidates:
        if i in used_rows or j in used_cols:
            continue
        matches.append((i, j))
        used_rows.add(i)
        used_cols.add(j)
    return matches


# ---------------------------------------------------------------------------
# ByteTracker
# ---------------------------------------------------------------------------


class ByteTracker:
    """简化版 ByteTrack 跟踪器。

    Args:
        track_thresh: 跟踪基础置信度阈值，低于此值不参与跟踪
        high_thresh: 高置信度阈值，分桶用
        match_thresh: IOU 匹配阈值
        max_lost_frames: 跟踪丢失最大帧数，超过移除
        min_box_area: 最小框面积，低于此值忽略
    """

    def __init__(
        self,
        track_thresh: float = 0.5,
        high_thresh: float = 0.6,
        match_thresh: float = 0.8,
        max_lost_frames: int = 30,
        min_box_area: float = 10.0,
    ) -> None:
        self.track_thresh = track_thresh
        self.high_thresh = high_thresh
        self.match_thresh = match_thresh
        self.max_lost_frames = max_lost_frames
        self.min_box_area = min_box_area

        self._tracks: Dict[int, Track] = {}
        self._next_id: int = 1
        self._frame_count: int = 0

    # ------------------------------------------------------------------
    # 主入口
    # ------------------------------------------------------------------

    def update(
        self,
        detections: List[UnifiedDetection],
    ) -> List[UnifiedDetection]:
        """对一帧检测结果执行跟踪更新，填充 track_id。

        Returns:
            带 track_id 的检测结果 (顺序保持不变)
        """

        self._frame_count += 1

        # 1. 过滤无效检测 (面积太小)
        valid: List[Tuple[int, UnifiedDetection]] = []
        for idx, det in enumerate(detections):
            if det.bbox.area < self.min_box_area:
                continue
            valid.append((idx, det))

        # 2. 按置信度分桶
        high_dets: List[Tuple[int, UnifiedDetection]] = []
        low_dets: List[Tuple[int, UnifiedDetection]] = []
        for idx, det in valid:
            if det.confidence >= self.high_thresh:
                high_dets.append((idx, det))
            elif det.confidence >= self.track_thresh:
                low_dets.append((idx, det))
            # < track_thresh 直接忽略

        # 3. 获取当前 tracked + lost 的 tracks 列表
        active_tracks: List[Track] = [
            t for t in self._tracks.values()
            if t.state in (TrackState.TRACKED, TrackState.LOST, TrackState.NEW)
        ]

        # 4. 第一轮: 与 high 检测匹配
        matches_high, unmatched_tracks, unmatched_high = self._match(
            active_tracks, [d for _, d in high_dets], self.match_thresh
        )
        for track_idx, det_idx in matches_high:
            track = active_tracks[track_idx]
            orig_idx, det = high_dets[det_idx]
            self._update_track(track, det)
            det.track_id = track.track_id

        # 5. 第二轮: 未匹配 tracks 与 low 检测匹配
        remaining_tracks = [active_tracks[i] for i in unmatched_tracks]
        matches_low, still_unmatched, unmatched_low = self._match(
            remaining_tracks, [d for _, d in low_dets], 0.5  # 低置信度用更宽松阈值
        )
        for track_idx, det_idx in matches_low:
            track = remaining_tracks[track_idx]
            orig_idx, det = low_dets[det_idx]
            self._update_track(track, det)
            det.track_id = track.track_id

        # 6. 未匹配的 high 检测创建新 track
        for det_idx in unmatched_high:
            orig_idx, det = high_dets[det_idx]
            new_track = self._create_track(det)
            det.track_id = new_track.track_id

        # 7. 未匹配的 tracks 增加 lost_frames
        unmatched_original_tracks = [remaining_tracks[i] for i in still_unmatched]
        for track in unmatched_original_tracks:
            track.lost_frames += 1
            track.state = TrackState.LOST
            if track.lost_frames > self.max_lost_frames:
                track.state = TrackState.REMOVED

        # 8. 清理已移除的 tracks
        self._tracks = {
            tid: t for tid, t in self._tracks.items()
            if t.state != TrackState.REMOVED
        }

        return detections

    # ------------------------------------------------------------------
    # 内部
    # ------------------------------------------------------------------

    def _match(
        self,
        tracks: List[Track],
        detections: List[UnifiedDetection],
        threshold: float,
    ) -> Tuple[List[Tuple[int, int]], List[int], List[int]]:
        """计算 tracks 与 detections 的 IOU 匹配。

        Returns:
            (matches, unmatched_track_indices, unmatched_det_indices)
        """

        if not tracks or not detections:
            return [], list(range(len(tracks))), list(range(len(detections)))

        # 构建 IOU 矩阵
        iou_matrix: List[List[float]] = []
        for t in tracks:
            row = [compute_iou(t.bbox, d.bbox) for d in detections]
            iou_matrix.append(row)

        matches = greedy_match(iou_matrix, threshold)
        matched_tracks = {i for i, _ in matches}
        matched_dets = {j for _, j in matches}
        unmatched_tracks = [i for i in range(len(tracks)) if i not in matched_tracks]
        unmatched_dets = [j for j in range(len(detections)) if j not in matched_dets]
        return matches, unmatched_tracks, unmatched_dets

    def _create_track(self, det: UnifiedDetection) -> Track:
        track_id = self._next_id
        self._next_id += 1
        track = Track(
            track_id=track_id,
            bbox=det.bbox,
            confidence=det.confidence,
            class_name=det.class_name,
            state=TrackState.TRACKED,
            age=1,
            lost_frames=0,
            last_update_frame=self._frame_count,
        )
        self._tracks[track_id] = track
        return track

    def _update_track(self, track: Track, det: UnifiedDetection) -> None:
        track.bbox = det.bbox
        track.confidence = det.confidence
        track.age += 1
        track.lost_frames = 0
        track.state = TrackState.TRACKED
        track.last_update_frame = self._frame_count

    # ------------------------------------------------------------------
    # 状态
    # ------------------------------------------------------------------

    @property
    def active_tracks(self) -> List[Track]:
        return [t for t in self._tracks.values() if t.state == TrackState.TRACKED]

    @property
    def lost_tracks(self) -> List[Track]:
        return [t for t in self._tracks.values() if t.state == TrackState.LOST]

    def reset(self) -> None:
        self._tracks.clear()
        self._next_id = 1
        self._frame_count = 0


# ---------------------------------------------------------------------------
# TrackingService: 多流跟踪器管理
# ---------------------------------------------------------------------------


class TrackingService:
    """多流目标跟踪服务。

    为每个 source_id (摄像头/视频流) 维护独立的 ByteTracker 实例。
    """

    def __init__(
        self,
        track_thresh: float = 0.5,
        high_thresh: float = 0.6,
        match_thresh: float = 0.8,
        max_lost_frames: int = 30,
        min_box_area: float = 10.0,
    ) -> None:
        self.track_thresh = track_thresh
        self.high_thresh = high_thresh
        self.match_thresh = match_thresh
        self.max_lost_frames = max_lost_frames
        self.min_box_area = min_box_area
        self._trackers: Dict[str, ByteTracker] = {}

    def get_or_create(self, source_id: str) -> ByteTracker:
        if source_id not in self._trackers:
            self._trackers[source_id] = ByteTracker(
                track_thresh=self.track_thresh,
                high_thresh=self.high_thresh,
                match_thresh=self.match_thresh,
                max_lost_frames=self.max_lost_frames,
                min_box_area=self.min_box_area,
            )
        return self._trackers[source_id]

    def update(
        self,
        source_id: str,
        detections: List[UnifiedDetection],
    ) -> List[UnifiedDetection]:
        tracker = self.get_or_create(source_id)
        return tracker.update(detections)

    def reset(self, source_id: Optional[str] = None) -> None:
        if source_id is None:
            self._trackers.clear()
        elif source_id in self._trackers:
            self._trackers[source_id].reset()

    def remove(self, source_id: str) -> None:
        self._trackers.pop(source_id, None)

    @property
    def source_count(self) -> int:
        return len(self._trackers)


__all__ = [
    "ByteTracker",
    "TrackingService",
    "Track",
    "TrackState",
    "compute_iou",
    "greedy_match",
]