import math
import time
from typing import List, Optional, Tuple

from saltybot_social.person_state import PersonTrack, State


class IdentityFuser:
    """Multi-modal identity fusion: face_id + speaker_id + UWB anchor -> unified person_id."""

    def __init__(self, position_window: float = 3.0):
        self._tracks: dict[int, PersonTrack] = {}
        self._next_id: int = 1
        self._position_window = position_window
        self._max_history = 20

    def _distance_3d(self, a: Optional[tuple], b: Optional[tuple]) -> float:
        if a is None or b is None:
            return float('inf')
        return math.sqrt(sum((ai - bi) ** 2 for ai, bi in zip(a, b)))

    def _compute_distance_and_bearing(self, position: tuple) -> Tuple[float, float]:
        x, y, z = position
        distance = math.sqrt(x * x + y * y + z * z)
        bearing = math.degrees(math.atan2(y, x))
        return distance, bearing

    def _allocate_id(self) -> int:
        pid = self._next_id
        self._next_id += 1
        return pid

    def update_face(self, face_id: int, name: str,
                    position_3d: Optional[tuple], camera_id: int,
                    now: float) -> int:
        """Match by face_id first. If face_id >= 0, find existing track or create new."""
        if face_id >= 0:
            for track in self._tracks.values():
                if track.face_id == face_id:
                    track.name = name if name and name != 'unknown' else track.name
                    if position_3d is not None:
                        track.position = position_3d
                        dist, bearing = self._compute_distance_and_bearing(position_3d)
                        track.distance = dist
                        track.bearing_deg = bearing
                        track.history_distances.append(dist)
                        if len(track.history_distances) > self._max_history:
                            track.history_distances = track.history_distances[-self._max_history:]
                    track.camera_id = camera_id
                    track.last_seen = now
                    track.last_face_seen = now
                    return track.person_id

        # No existing track with this face_id; try proximity match for face_id < 0
        if face_id < 0 and position_3d is not None:
            best_track = None
            best_dist = 0.5  # 0.5m proximity threshold
            for track in self._tracks.values():
                d = self._distance_3d(track.position, position_3d)
                if d < best_dist and (now - track.last_seen) < self._position_window:
                    best_dist = d
                    best_track = track
            if best_track is not None:
                best_track.position = position_3d
                dist, bearing = self._compute_distance_and_bearing(position_3d)
                best_track.distance = dist
                best_track.bearing_deg = bearing
                best_track.history_distances.append(dist)
                if len(best_track.history_distances) > self._max_history:
                    best_track.history_distances = best_track.history_distances[-self._max_history:]
                best_track.camera_id = camera_id
                best_track.last_seen = now
                return best_track.person_id

        # Create new track
        pid = self._allocate_id()
        track = PersonTrack(person_id=pid, face_id=face_id, name=name,
                            camera_id=camera_id, last_seen=now, last_face_seen=now)
        if position_3d is not None:
            track.position = position_3d
            dist, bearing = self._compute_distance_and_bearing(position_3d)
            track.distance = dist
            track.bearing_deg = bearing
            track.history_distances.append(dist)
        self._tracks[pid] = track
        return pid

    def update_speaker(self, speaker_id: str, now: float) -> int:
        """Find track with nearest recently-seen position, assign speaker_id."""
        # First check if any track already has this speaker_id
        for track in self._tracks.values():
            if track.speaker_id == speaker_id:
                track.last_seen = now
                return track.person_id

        # Assign to nearest recently-seen track
        best_track = None
        best_dist = float('inf')
        for track in self._tracks.values():
            if (now - track.last_seen) < self._position_window and track.distance > 0:
                if track.distance < best_dist:
                    best_dist = track.distance
                    best_track = track

        if best_track is not None:
            best_track.speaker_id = speaker_id
            best_track.last_seen = now
            return best_track.person_id

        # No suitable track found; create a new one
        pid = self._allocate_id()
        track = PersonTrack(person_id=pid, speaker_id=speaker_id,
                            last_seen=now)
        self._tracks[pid] = track
        return pid

    def update_uwb(self, uwb_anchor_id: str, position_3d: tuple,
                    now: float) -> int:
        """Match by proximity (nearest track within 0.5m)."""
        # Check existing UWB assignment
        for track in self._tracks.values():
            if track.uwb_anchor_id == uwb_anchor_id:
                track.position = position_3d
                dist, bearing = self._compute_distance_and_bearing(position_3d)
                track.distance = dist
                track.bearing_deg = bearing
                track.history_distances.append(dist)
                if len(track.history_distances) > self._max_history:
                    track.history_distances = track.history_distances[-self._max_history:]
                track.last_seen = now
                return track.person_id

        # Proximity match
        best_track = None
        best_dist = 0.5
        for track in self._tracks.values():
            d = self._distance_3d(track.position, position_3d)
            if d < best_dist and (now - track.last_seen) < self._position_window:
                best_dist = d
                best_track = track

        if best_track is not None:
            best_track.uwb_anchor_id = uwb_anchor_id
            best_track.position = position_3d
            dist, bearing = self._compute_distance_and_bearing(position_3d)
            best_track.distance = dist
            best_track.bearing_deg = bearing
            best_track.history_distances.append(dist)
            if len(best_track.history_distances) > self._max_history:
                best_track.history_distances = best_track.history_distances[-self._max_history:]
            best_track.last_seen = now
            return best_track.person_id

        # Create new track
        pid = self._allocate_id()
        track = PersonTrack(person_id=pid, uwb_anchor_id=uwb_anchor_id,
                            position=position_3d, last_seen=now)
        dist, bearing = self._compute_distance_and_bearing(position_3d)
        track.distance = dist
        track.bearing_deg = bearing
        track.history_distances.append(dist)
        self._tracks[pid] = track
        return pid

    def get_all_tracks(self) -> List[PersonTrack]:
        return list(self._tracks.values())

    @staticmethod
    def compute_attention(tracks: List[PersonTrack]) -> int:
        """Focus on nearest engaged/talking person, or -1 if none."""
        candidates = [t for t in tracks
                      if t.state in (State.ENGAGED, State.TALKING) and t.distance > 0]
        if not candidates:
            return -1
        # Prefer talking, then nearest engaged
        talking = [t for t in candidates if t.state == State.TALKING]
        if talking:
            return min(talking, key=lambda t: t.distance).person_id
        return min(candidates, key=lambda t: t.distance).person_id

    def prune_absent(self, absent_timeout: float = 30.0):
        """Remove tracks absent longer than timeout."""
        now = time.monotonic()
        to_remove = [pid for pid, t in self._tracks.items()
                     if t.state == State.ABSENT and (now - t.last_seen) > absent_timeout]
        for pid in to_remove:
            del self._tracks[pid]

    @staticmethod
    def detect_group(tracks: List[PersonTrack]) -> bool:
        """Returns True if >= 2 persons within 1.5m of each other."""
        active = [t for t in tracks
                  if t.position is not None and t.state != State.ABSENT]
        for i, a in enumerate(active):
            for b in active[i + 1:]:
                dx = a.position[0] - b.position[0]
                dy = a.position[1] - b.position[1]
                dz = a.position[2] - b.position[2]
                if math.sqrt(dx * dx + dy * dy + dz * dz) < 1.5:
                    return True
        return False
