"""
Object tracking management and state handling.

This module provides tracking state management, session handling,
and occupancy detection functionality for the detection pipeline.
"""

import time
import logging
from typing import Dict, List, Any, Optional, Tuple, Set
from dataclasses import dataclass, field

from ..core.constants import SESSION_TIMEOUT_SECONDS

logger = logging.getLogger(__name__)


@dataclass
class SessionState:
    """Session state for tracking management."""
    active: bool = True
    waiting_for_backend_session: bool = False
    wait_start_time: float = 0.0
    reset_tracker_on_resume: bool = False
    occupancy_mode: bool = False
    occupancy_enabled_at: Optional[float] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format."""
        return {
            "active": self.active,
            "waiting_for_backend_session": self.waiting_for_backend_session,
            "wait_start_time": self.wait_start_time,
            "reset_tracker_on_resume": self.reset_tracker_on_resume,
            "occupancy_mode": self.occupancy_mode,
            "occupancy_enabled_at": self.occupancy_enabled_at
        }
    
    def from_dict(self, data: Dict[str, Any]) -> None:
        """Update from dictionary data."""
        self.active = data.get("active", True)
        self.waiting_for_backend_session = data.get("waiting_for_backend_session", False)
        self.wait_start_time = data.get("wait_start_time", 0.0)
        self.reset_tracker_on_resume = data.get("reset_tracker_on_resume", False)
        self.occupancy_mode = data.get("occupancy_mode", False)
        self.occupancy_enabled_at = data.get("occupancy_enabled_at")


@dataclass
class TrackingData:
    """Complete tracking data for a camera and model."""
    track_stability_counters: Dict[int, int] = field(default_factory=dict)
    stable_tracks: Set[int] = field(default_factory=set)
    session_state: SessionState = field(default_factory=SessionState)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format."""
        return {
            "track_stability_counters": dict(self.track_stability_counters),
            "stable_tracks": list(self.stable_tracks),
            "session_state": self.session_state.to_dict()
        }


class TrackingManager:
    """
    Manages object tracking state and session handling across cameras and models.
    
    This class provides centralized tracking state management including:
    - Track stability counters
    - Session state management
    - Backend session timeout handling
    - Occupancy detection mode
    """
    
    def __init__(self, session_timeout_seconds: int = SESSION_TIMEOUT_SECONDS):
        """
        Initialize tracking manager.
        
        Args:
            session_timeout_seconds: Timeout for backend session waiting
        """
        self.session_timeout_seconds = session_timeout_seconds
        self._camera_tracking_data: Dict[str, Dict[str, TrackingData]] = {}
        self._lock = None
    
    def _ensure_thread_safety(self):
        """Initialize thread safety if not already done."""
        if self._lock is None:
            import threading
            self._lock = threading.RLock()
    
    def get_camera_tracking_data(self, camera_id: str, model_id: str) -> TrackingData:
        """
        Get or create tracking data for a specific camera and model.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            Tracking data for the camera and model
        """
        self._ensure_thread_safety()
        
        with self._lock:
            if camera_id not in self._camera_tracking_data:
                self._camera_tracking_data[camera_id] = {}
            
            if model_id not in self._camera_tracking_data[camera_id]:
                logger.warning(f"🔄 Camera {camera_id}: Creating NEW tracking data for {model_id} - this will reset any cooldown!")
                self._camera_tracking_data[camera_id][model_id] = TrackingData()
            
            return self._camera_tracking_data[camera_id][model_id]
    
    def check_stable_tracks(self, 
                          camera_id: str, 
                          model_id: str, 
                          regions_dict: Dict[str, Any]) -> Tuple[bool, List[Dict[str, Any]]]:
        """
        Check if any stable tracks match the detected classes for a specific camera.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            regions_dict: Dictionary mapping class names to detection regions
            
        Returns:
            Tuple of (has_stable_tracks, stable_detections)
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            stable_tracks = tracking_data.stable_tracks
            
            if not stable_tracks:
                return False, []
            
            # Check for track-based stability
            stable_detections = []
            
            for class_name, region_data in regions_dict.items():
                detection = region_data.get("detection", {})
                track_id = detection.get("id")
                
                if track_id in stable_tracks:
                    stable_detections.append(detection)
            
            has_stable = len(stable_detections) > 0
            logger.debug(f"Camera {camera_id}: Stable track check - stable_tracks: {list(stable_tracks)}, found: {has_stable}")
            
            return has_stable, stable_detections
    
    def reset_tracking_state(self, 
                           camera_id: str, 
                           model_id: str, 
                           reason: str = "session ended") -> None:
        """
        Reset tracking state after session completion or timeout.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            reason: Reason for reset (for logging)
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            # Clear all tracking data for fresh start
            tracking_data.track_stability_counters.clear()
            tracking_data.stable_tracks.clear()
            session_state.active = True
            session_state.waiting_for_backend_session = False
            session_state.wait_start_time = 0.0
            session_state.reset_tracker_on_resume = True
            
            logger.info(f"Camera {camera_id}: 🔄 Reset tracking state - {reason}")
            logger.info(f"Camera {camera_id}: 🧹 Cleared stability counters and stable tracks for fresh session")
    
    def is_camera_active(self, camera_id: str, model_id: str) -> bool:
        """
        Check if camera should be processing detections.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            True if camera should process detections, False otherwise
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            # Check if waiting for backend sessionId has timed out
            if session_state.waiting_for_backend_session:
                current_time = time.time()
                wait_start_time = session_state.wait_start_time
                elapsed_time = current_time - wait_start_time
                
                if elapsed_time >= self.session_timeout_seconds:
                    logger.warning(f"Camera {camera_id}: Backend sessionId timeout ({self.session_timeout_seconds}s) - resetting tracking")
                    self.reset_tracking_state(camera_id, model_id, "backend sessionId timeout")
                    return True
                else:
                    remaining_time = self.session_timeout_seconds - elapsed_time
                    logger.debug(f"Camera {camera_id}: Still waiting for backend sessionId ({remaining_time:.1f}s remaining)")
                    return False
            
            return session_state.active
    
    def set_waiting_for_backend_session(self, 
                                      camera_id: str, 
                                      model_id: str, 
                                      waiting: bool = True) -> None:
        """
        Set whether camera is waiting for backend session ID.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            waiting: Whether to wait for backend session
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            session_state.waiting_for_backend_session = waiting
            if waiting:
                session_state.wait_start_time = time.time()
                logger.debug(f"Camera {camera_id}: Started waiting for backend sessionId")
            else:
                session_state.wait_start_time = 0.0
                logger.debug(f"Camera {camera_id}: Stopped waiting for backend sessionId")
    
    def cleanup_camera_tracking(self, camera_id: str) -> None:
        """
        Clean up tracking data when a camera is disconnected.
        
        Args:
            camera_id: Unique camera identifier
        """
        self._ensure_thread_safety()
        
        with self._lock:
            if camera_id in self._camera_tracking_data:
                del self._camera_tracking_data[camera_id]
                logger.info(f"Cleaned up tracking data for camera {camera_id}")
    
    def set_occupancy_mode(self, 
                         camera_id: str, 
                         model_id: str, 
                         enable: bool = True) -> bool:
        """
        Enable or disable occupancy detection mode.
        
        Occupancy mode stops model inference after pipeline completion
        while backend session handling continues in background.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            enable: True to enable occupancy mode, False to disable
            
        Returns:
            Current occupancy mode state
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            if enable:
                session_state.occupancy_mode = True
                session_state.occupancy_enabled_at = time.time()
                logger.debug(f"Camera {camera_id}: Occupancy mode ENABLED - model will stop after pipeline completion")
            else:
                session_state.occupancy_mode = False
                session_state.occupancy_enabled_at = None
                logger.debug(f"Camera {camera_id}: Occupancy mode DISABLED - model will continue running")
            
            return session_state.occupancy_mode
    
    def is_occupancy_mode_enabled(self, camera_id: str, model_id: str) -> bool:
        """
        Check if occupancy mode is enabled for a camera.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            True if occupancy mode is enabled
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            return tracking_data.session_state.occupancy_mode
    
    def get_occupancy_duration(self, camera_id: str, model_id: str) -> Optional[float]:
        """
        Get duration since occupancy mode was enabled.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            Duration in seconds since occupancy mode was enabled, or None if not enabled
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            if not session_state.occupancy_mode or not session_state.occupancy_enabled_at:
                return None
            
            return time.time() - session_state.occupancy_enabled_at
    
    def get_session_state(self, camera_id: str, model_id: str) -> SessionState:
        """
        Get session state for a camera and model.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            Session state object
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            return tracking_data.session_state
    
    def update_session_state(self, 
                           camera_id: str, 
                           model_id: str, 
                           **kwargs) -> None:
        """
        Update session state properties.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            **kwargs: Session state properties to update
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            session_state = tracking_data.session_state
            
            for key, value in kwargs.items():
                if hasattr(session_state, key):
                    setattr(session_state, key, value)
                    logger.debug(f"Camera {camera_id}: Updated session state {key} = {value}")
                else:
                    logger.warning(f"Camera {camera_id}: Unknown session state property: {key}")
    
    def get_tracking_statistics(self, camera_id: str, model_id: str) -> Dict[str, Any]:
        """
        Get comprehensive tracking statistics for a camera and model.
        
        Args:
            camera_id: Unique camera identifier
            model_id: Model identifier
            
        Returns:
            Dictionary with tracking statistics
        """
        self._ensure_thread_safety()
        
        with self._lock:
            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
            current_time = time.time()
            
            stats = tracking_data.to_dict()
            
            # Add computed statistics
            stats["total_tracked_objects"] = len(tracking_data.track_stability_counters)
            stats["stable_track_count"] = len(tracking_data.stable_tracks)
            
            session_state = tracking_data.session_state
            if session_state.waiting_for_backend_session and session_state.wait_start_time > 0:
                stats["backend_wait_duration"] = current_time - session_state.wait_start_time
                stats["backend_wait_remaining"] = max(0, self.session_timeout_seconds - stats["backend_wait_duration"])
            
            if session_state.occupancy_mode and session_state.occupancy_enabled_at:
                stats["occupancy_duration"] = current_time - session_state.occupancy_enabled_at
            
            return stats
    
    def get_all_camera_stats(self) -> Dict[str, Dict[str, Dict[str, Any]]]:
        """
        Get tracking statistics for all monitored cameras.
        
        Returns:
            Nested dictionary: {camera_id: {model_id: stats}}
        """
        self._ensure_thread_safety()
        
        with self._lock:
            all_stats = {}
            for camera_id, models in self._camera_tracking_data.items():
                all_stats[camera_id] = {}
                for model_id in models.keys():
                    all_stats[camera_id][model_id] = self.get_tracking_statistics(camera_id, model_id)
            
            return all_stats


# Global tracking manager instance
tracking_manager = TrackingManager()


# ===== CONVENIENCE FUNCTIONS =====
# These provide the same interface as the original functions in pympta.py

def check_stable_tracks(camera_id: str, model_id: str, regions_dict: Dict[str, Any]) -> Tuple[bool, List[Dict[str, Any]]]:
    """Check if any stable tracks match the detected classes for a specific camera."""
    return tracking_manager.check_stable_tracks(camera_id, model_id, regions_dict)


def reset_tracking_state(camera_id: str, model_id: str, reason: str = "session ended") -> None:
    """Reset tracking state after session completion or timeout."""
    tracking_manager.reset_tracking_state(camera_id, model_id, reason)


def is_camera_active(camera_id: str, model_id: str) -> bool:
    """Check if camera should be processing detections."""
    return tracking_manager.is_camera_active(camera_id, model_id)


def cleanup_camera_stability(camera_id: str) -> None:
    """Clean up tracking data when a camera is disconnected."""
    tracking_manager.cleanup_camera_tracking(camera_id)


def occupancy_detector(camera_id: str, model_id: str, enable: bool = True) -> bool:
    """
    Enable or disable occupancy detection mode.
    
    When enabled:
    - Model stops inference after completing full pipeline
    - Backend sessionId handling continues in background
    
    Note: This is a temporary function that will be changed in the future.
    """
    return tracking_manager.set_occupancy_mode(camera_id, model_id, enable)


def get_camera_tracking_data(camera_id: str, model_id: str) -> Dict[str, Any]:
    """Get tracking data for a specific camera and model."""
    tracking_data = tracking_manager.get_camera_tracking_data(camera_id, model_id)
    return tracking_data.to_dict()


def set_waiting_for_backend_session(camera_id: str, model_id: str, waiting: bool = True) -> None:
    """Set whether camera is waiting for backend session ID."""
    tracking_manager.set_waiting_for_backend_session(camera_id, model_id, waiting)


def get_tracking_statistics(camera_id: str, model_id: str) -> Dict[str, Any]:
    """Get comprehensive tracking statistics for a camera and model."""
    return tracking_manager.get_tracking_statistics(camera_id, model_id)