"""Multi-person tracker for Issue #423.

Tracks up to 10 people with:
- Unique IDs (persistent across frames)
- Target priority: wake-word speaker > closest known > largest bbox
- Occlusion handoff (3s grace period)
- Re-ID via face embedding + HSV color histogram
- Group detection (follow centroid)
- Lost target: stop + rotate + SEARCHING state
"""

from __future__ import annotations

import time
import numpy as np
import cv2
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple


@dataclass
class TrackedPerson:
    """A tracked person with identity and state."""
    track_id: int
    bearing_rad: float
    distance_m: float
    confidence: float
    is_speaking: bool = False
    source: str = "camera"
    last_seen_time: float = field(default_factory=time.time)
    embedding: Optional[np.ndarray] = None  # Face embedding for re-ID
    hsv_histogram: Optional[np.ndarray] = None  # HSV histogram fallback
    bbox: Optional[Tuple[int, int, int, int]] = None  # (x, y, w, h)
    is_occluded: bool = False
    occlusion_start_time: float = 0.0

    def age_seconds(self) -> float:
        """Time since last detection."""
        return time.time() - self.last_seen_time

    def mark_seen(
        self,
        bearing: float,
        distance: float,
        confidence: float,
        is_speaking: bool = False,
        embedding: Optional[np.ndarray] = None,
        hsv_hist: Optional[np.ndarray] = None,
        bbox: Optional[Tuple] = None,
    ) -> None:
        """Update person state after detection."""
        self.bearing_rad = bearing
        self.distance_m = distance
        self.confidence = confidence
        self.is_speaking = is_speaking
        self.last_seen_time = time.time()
        self.is_occluded = False
        if embedding is not None:
            self.embedding = embedding
        if hsv_hist is not None:
            self.hsv_histogram = hsv_hist
        if bbox is not None:
            self.bbox = bbox

    def mark_occluded(self) -> None:
        """Mark person as occluded but within grace period."""
        if not self.is_occluded:
            self.is_occluded = True
            self.occlusion_start_time = time.time()


class MultiPersonTracker:
    """Tracks multiple people with re-ID and group detection."""

    def __init__(
        self,
        max_people: int = 10,
        occlusion_grace_s: float = 3.0,
        embedding_similarity_threshold: float = 0.75,
        hsv_similarity_threshold: float = 0.60,
    ):
        self.max_people = max_people
        self.occlusion_grace_s = occlusion_grace_s
        self.embedding_threshold = embedding_similarity_threshold
        self.hsv_threshold = hsv_similarity_threshold

        self.tracked_people: Dict[int, TrackedPerson] = {}
        self.next_track_id = 0
        self.active_target_id: Optional[int] = None
        self.searching: bool = False
        self.search_bearing: float = 0.0

    # ── Main tracking update ────────────────────────────────────────────────

    def update(
        self,
        detections: List[Dict],
        speaking_person_id: Optional[int] = None,
    ) -> Tuple[List[TrackedPerson], Optional[int]]:
        """Update tracker with new detections.

        Args:
            detections: List of dicts with keys:
                - bearing_rad, distance_m, confidence
                - embedding (optional): face embedding vector
                - hsv_histogram (optional): HSV color histogram
                - bbox (optional): (x, y, w, h)
            speaking_person_id: wake-word speaker (highest priority)

        Returns:
            (all_tracked_people, active_target_id)
        """
        # Clean up stale tracks
        self._prune_lost_people()

        # Match detections to existing tracks
        matched_ids, unmatched_detections = self._match_detections(detections)

        # Update matched tracks
        for track_id, det_idx in matched_ids:
            det = detections[det_idx]
            self.tracked_people[track_id].mark_seen(
                bearing=det["bearing_rad"],
                distance=det["distance_m"],
                confidence=det["confidence"],
                is_speaking=det.get("is_speaking", False),
                embedding=det.get("embedding"),
                hsv_hist=det.get("hsv_histogram"),
                bbox=det.get("bbox"),
            )

        # Create new tracks for unmatched detections
        for det_idx in unmatched_detections:
            if len(self.tracked_people) < self.max_people:
                det = detections[det_idx]
                track_id = self.next_track_id
                self.next_track_id += 1
                self.tracked_people[track_id] = TrackedPerson(
                    track_id=track_id,
                    bearing_rad=det["bearing_rad"],
                    distance_m=det["distance_m"],
                    confidence=det["confidence"],
                    is_speaking=det.get("is_speaking", False),
                    embedding=det.get("embedding"),
                    hsv_histogram=det.get("hsv_histogram"),
                    bbox=det.get("bbox"),
                )

        # Update target priority
        self._update_target_priority(speaking_person_id)

        return list(self.tracked_people.values()), self.active_target_id

    # ── Detection-to-track matching ──────────────────────────────────────────

    def _match_detections(
        self, detections: List[Dict]
    ) -> Tuple[List[Tuple[int, int]], List[int]]:
        """Match detections to tracked people via re-ID.

        Returns:
            (matched_pairs, unmatched_detection_indices)
        """
        if not detections or not self.tracked_people:
            return [], list(range(len(detections)))

        matched = []
        used_det_indices = set()

        # Re-ID matching using embeddings and HSV histograms
        for track_id, person in self.tracked_people.items():
            best_score = -1.0
            best_det_idx = -1

            for det_idx, det in enumerate(detections):
                if det_idx in used_det_indices:
                    continue

                score = self._compute_similarity(person, det)
                if score > best_score:
                    best_score = score
                    best_det_idx = det_idx

            # Match if above threshold
            if best_det_idx >= 0:
                if (
                    (person.embedding is not None and best_score >= self.embedding_threshold)
                    or (person.hsv_histogram is not None and best_score >= self.hsv_threshold)
                ):
                    matched.append((track_id, best_det_idx))
                    used_det_indices.add(best_det_idx)
                    person.is_occluded = False

        unmatched_det_indices = [i for i in range(len(detections)) if i not in used_det_indices]
        return matched, unmatched_det_indices

    def _compute_similarity(self, person: TrackedPerson, detection: Dict) -> float:
        """Compute similarity between tracked person and detection."""
        if person.embedding is not None and detection.get("embedding") is not None:
            # Cosine similarity via dot product (embeddings are L2-normalized)
            emb_score = float(np.dot(person.embedding, detection["embedding"]))
            return emb_score

        if person.hsv_histogram is not None and detection.get("hsv_histogram") is not None:
            # Compare HSV histograms
            hist_score = float(
                cv2.compareHist(
                    person.hsv_histogram,
                    detection["hsv_histogram"],
                    cv2.HISTCMP_COSINE,
                )
            )
            return hist_score

        # Fallback: spatial proximity (within similar distance/bearing)
        bearing_diff = abs(person.bearing_rad - detection["bearing_rad"])
        distance_diff = abs(person.distance_m - detection["distance_m"])
        spatial_score = 1.0 / (1.0 + bearing_diff + distance_diff * 0.1)
        return spatial_score

    # ── Target priority and group handling ────────────────────────────────────

    def _update_target_priority(self, speaking_person_id: Optional[int]) -> None:
        """Update active target based on priority:
        1. Wake-word speaker
        2. Closest known person
        3. Largest bounding box
        """
        if not self.tracked_people:
            self.active_target_id = None
            self.searching = True
            return

        # Priority 1: Wake-word speaker
        if speaking_person_id is not None and speaking_person_id in self.tracked_people:
            self.active_target_id = speaking_person_id
            self.searching = False
            return

        # Priority 2: Closest non-occluded person
        closest_id = None
        closest_dist = float("inf")
        for track_id, person in self.tracked_people.items():
            if not person.is_occluded and person.distance_m > 0:
                if person.distance_m < closest_dist:
                    closest_dist = person.distance_m
                    closest_id = track_id

        if closest_id is not None:
            self.active_target_id = closest_id
            self.searching = False
            return

        # Priority 3: Largest bounding box
        largest_id = None
        largest_area = 0
        for track_id, person in self.tracked_people.items():
            if person.bbox is not None:
                _, _, w, h = person.bbox
                area = w * h
                if area > largest_area:
                    largest_area = area
                    largest_id = track_id

        if largest_id is not None:
            self.active_target_id = largest_id
            self.searching = False
        else:
            self.searching = True

    def get_group_centroid(self) -> Optional[Tuple[float, float, float]]:
        """Get centroid (bearing, distance) of all tracked people.

        Returns:
            (mean_bearing_rad, mean_distance_m) or None if no people
        """
        if not self.tracked_people:
            return None

        bearings = []
        distances = []
        for person in self.tracked_people.values():
            if not person.is_occluded and person.distance_m > 0:
                bearings.append(person.bearing_rad)
                distances.append(person.distance_m)

        if not bearings:
            return None

        # Mean bearing (circular mean)
        mean_bearing = float(np.mean(bearings))
        mean_distance = float(np.mean(distances))
        return mean_bearing, mean_distance

    # ── Pruning and cleanup ──────────────────────────────────────────────────

    def _prune_lost_people(self) -> None:
        """Remove people lost for more than occlusion grace period."""
        current_time = time.time()
        to_remove = []

        for track_id, person in self.tracked_people.items():
            # Check occlusion timeout
            if person.is_occluded:
                elapsed = current_time - person.occlusion_start_time
                if elapsed > self.occlusion_grace_s:
                    to_remove.append(track_id)
                    continue

            # Check long-term invisibility (10s)
            if person.age_seconds() > 10.0:
                to_remove.append(track_id)

        for track_id in to_remove:
            del self.tracked_people[track_id]
            if self.active_target_id == track_id:
                self.active_target_id = None
                self.searching = True
