"""
HTTP Streaming Session Manager
==============================

Thin wrapper around BADAS.predict_frame() for HTTP streaming.

Clients POST individual frames (JPEG) and receive predictions when
the 16-frame buffer is full. All buffer management, preprocessing,
and inference is handled by the BADAS class.

Usage:
    1. POST /api/stream/session/start -> get session_id
    2. POST /api/stream/session/{id}/frame with JPEG data
       -> returns prediction when buffer full, otherwise {"status": "buffering"}
    3. DELETE /api/stream/session/{id} when done

Example:
    # Start session
    session = requests.post("/api/stream/session/start", json={"stride": 1})
    session_id = session.json()["session_id"]

    # Stream frames
    for frame_jpeg in camera_stream:
        result = requests.post(f"/api/stream/session/{session_id}/frame", data=frame_jpeg)
        if result.json().get("type") == "prediction":
            print(f"Probability: {result.json()['probability']}")

    # Cleanup
    requests.delete(f"/api/stream/session/{session_id}")
"""

import time
import uuid
import threading
import numpy as np
from typing import Dict, List, Optional, Any, TYPE_CHECKING
from dataclasses import dataclass
from threading import Lock
from io import BytesIO

import cv2

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

if TYPE_CHECKING:
    from badas.inference import BADAS


@dataclass
class SessionConfig:
    """Configuration for a streaming session."""
    smoothing_enabled: bool = True
    timeout_seconds: int = 300  # 5 minutes


@dataclass
class SessionStats:
    """Statistics for a streaming session."""
    session_id: str
    frame_count: int
    buffer_fill: int
    buffer_size: int
    predictions_made: int
    created_at: float
    last_activity: float
    avg_latency_ms: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "session_id": self.session_id,
            "frame_count": self.frame_count,
            "buffer_fill": self.buffer_fill,
            "buffer_size": self.buffer_size,
            "predictions_made": self.predictions_made,
            "uptime_seconds": round(time.time() - self.created_at, 1),
            "last_activity_seconds_ago": round(time.time() - self.last_activity, 1),
            "avg_latency_ms": round(self.avg_latency_ms, 1),
        }


class StreamingSession:
    """
    Manages a single HTTP streaming session.

    This is a thin wrapper around BADAS.predict_frame() that:
    - Decodes JPEG frames
    - Tracks session statistics
    - Handles session lifecycle

    All buffer management, preprocessing, and inference is delegated to BADAS.
    """

    def __init__(
        self,
        session_id: str,
        predictor: "BADAS",
        config: SessionConfig,
    ):
        self.session_id = session_id
        self.predictor = predictor
        self.config = config

        # Session tracking
        self.frame_count = 0
        self.predictions_made = 0
        self.created_at = time.time()
        self.last_activity = time.time()
        self.latencies: List[float] = []

        # Thread safety
        self.lock = Lock()

        # Reset predictor state for this session
        self.predictor.reset_streaming_state()

    def add_frame(self, frame_data: bytes) -> Dict[str, Any]:
        """
        Add a JPEG-encoded frame to the session.

        Args:
            frame_data: JPEG image bytes

        Returns:
            Result from BADAS.predict_frame() with session_id added
        """
        with self.lock:
            self.last_activity = time.time()

            try:
                # Decode JPEG to numpy array (BGR for OpenCV)
                frame = self._decode_frame(frame_data)
                if frame is None:
                    return {"error": "Failed to decode frame", "session_id": self.session_id}

                self.frame_count += 1

                # Delegate to BADAS.predict_frame() - all logic is there
                result = self.predictor.predict_frame(
                    frame,
                    apply_smoothing=self.config.smoothing_enabled
                )

                # Add session_id to result
                result["session_id"] = self.session_id

                # Track statistics
                if result.get("type") == "prediction":
                    self.predictions_made += 1
                    if "latency_ms" in result:
                        self.latencies.append(result["latency_ms"])
                        if len(self.latencies) > 100:
                            self.latencies.pop(0)

                return result

            except Exception as e:
                return {"error": f"Frame processing error: {str(e)}", "session_id": self.session_id}

    def _decode_frame(self, frame_data: bytes) -> Optional[np.ndarray]:
        """Decode JPEG bytes to BGR numpy array."""
        try:
            # Try OpenCV first (faster)
            nparr = np.frombuffer(frame_data, np.uint8)
            frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            if frame is not None:
                return frame  # Already BGR

            # Fallback to PIL
            if PIL_AVAILABLE:
                img = Image.open(BytesIO(frame_data))
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                # Convert RGB to BGR for OpenCV compatibility
                return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

            return None

        except Exception:
            return None

    def is_expired(self) -> bool:
        """Check if session has expired due to inactivity."""
        return (time.time() - self.last_activity) > self.config.timeout_seconds

    def get_stats(self) -> SessionStats:
        """Get session statistics."""
        with self.lock:
            avg_latency = sum(self.latencies) / len(self.latencies) if self.latencies else 0.0

            # Get buffer info from predictor
            buffer_fill = 0
            buffer_size = getattr(self.predictor, 'window_size', 16)
            if hasattr(self.predictor, '_gpu_preprocessor') and self.predictor._gpu_preprocessor:
                buffer_fill = getattr(self.predictor._gpu_preprocessor, 'buffer_count', 0)

            return SessionStats(
                session_id=self.session_id,
                frame_count=self.frame_count,
                buffer_fill=buffer_fill,
                buffer_size=buffer_size,
                predictions_made=self.predictions_made,
                created_at=self.created_at,
                last_activity=self.last_activity,
                avg_latency_ms=avg_latency,
            )

    def reset(self):
        """Reset session state for reuse."""
        with self.lock:
            self.frame_count = 0
            self.predictions_made = 0
            self.latencies.clear()
            self.predictor.reset_streaming_state()


class StreamingSessionManager:
    """
    Manages multiple HTTP streaming sessions.

    Handles session lifecycle, cleanup of expired sessions, and
    shared predictor management.
    """

    def __init__(self, predictor: "BADAS" = None):
        """
        Initialize session manager.

        Args:
            predictor: BADAS instance to use for all sessions.
                      If None, must be set later via set_predictor().
        """
        self._predictor = predictor
        self._sessions: Dict[str, StreamingSession] = {}
        self._lock = Lock()
        self._cleanup_thread: Optional[threading.Thread] = None
        self._running = False

    def set_predictor(self, predictor: "BADAS"):
        """Set the predictor instance for new sessions."""
        self._predictor = predictor

    def start(self):
        """Start the session manager (begins cleanup thread)."""
        if self._running:
            return
        self._running = True
        self._cleanup_thread = threading.Thread(target=self._cleanup_worker, daemon=True)
        self._cleanup_thread.start()

    def stop(self):
        """Stop the session manager."""
        self._running = False
        with self._lock:
            self._sessions.clear()

    def create_session(
        self,
        stride: int = 1,  # Kept for API compatibility, but BADAS handles it
        smoothing_enabled: bool = True,
        resize_mode: Optional[str] = None,  # Kept for API compatibility
    ) -> str:
        """
        Create a new streaming session.

        Args:
            stride: (Ignored - BADAS uses stride=1 for predict_frame)
            smoothing_enabled: Apply temporal smoothing
            resize_mode: (Ignored - uses predictor's configured mode)

        Returns:
            session_id: Unique session identifier
        """
        if self._predictor is None:
            raise RuntimeError("Predictor not set. Call set_predictor() first.")

        session_id = str(uuid.uuid4())

        config = SessionConfig(
            smoothing_enabled=smoothing_enabled,
        )

        session = StreamingSession(session_id, self._predictor, config)

        with self._lock:
            self._sessions[session_id] = session

        return session_id

    def get_session(self, session_id: str) -> Optional[StreamingSession]:
        """Get a session by ID."""
        with self._lock:
            return self._sessions.get(session_id)

    def add_frame(self, session_id: str, frame_data: bytes) -> Dict[str, Any]:
        """
        Add a frame to a session.

        Args:
            session_id: Session identifier
            frame_data: JPEG image bytes

        Returns:
            Prediction result or buffering status
        """
        session = self.get_session(session_id)
        if session is None:
            return {"error": "Session not found"}

        return session.add_frame(frame_data)

    def delete_session(self, session_id: str) -> bool:
        """Delete a session."""
        with self._lock:
            if session_id in self._sessions:
                del self._sessions[session_id]
                return True
        return False

    def get_session_stats(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Get statistics for a specific session."""
        session = self.get_session(session_id)
        if session is None:
            return None
        return session.get_stats().to_dict()

    def get_all_stats(self) -> List[Dict[str, Any]]:
        """Get statistics for all active sessions."""
        with self._lock:
            return [s.get_stats().to_dict() for s in self._sessions.values()]

    def _cleanup_worker(self):
        """Background thread that cleans up expired sessions."""
        while self._running:
            time.sleep(60)  # Check every minute
            try:
                self._cleanup_expired()
            except Exception:
                pass

    def _cleanup_expired(self):
        """Remove expired sessions."""
        with self._lock:
            expired = [
                sid for sid, session in self._sessions.items()
                if session.is_expired()
            ]
            for sid in expired:
                del self._sessions[sid]


# Global session manager instance
_session_manager: Optional[StreamingSessionManager] = None


def get_session_manager() -> StreamingSessionManager:
    """Get or create the global session manager."""
    global _session_manager
    if _session_manager is None:
        _session_manager = StreamingSessionManager()
    return _session_manager


def init_session_manager(predictor: "BADAS") -> StreamingSessionManager:
    """Initialize the global session manager with a predictor."""
    global _session_manager
    _session_manager = StreamingSessionManager(predictor)
    _session_manager.start()
    return _session_manager
