"""
Streaming API module for Valence SDK.

This module provides WebSocket-based real-time emotion detection functionality.
Requires python-socketio[client] to be installed.
"""
import base64
try:
    import socketio
except ImportError:
    raise ImportError(
        "python-socketio is required for streaming support. "
        "Install it with: pip install 'python-socketio[client]>=5.10.0'"
    )

from .logger import get_logger
from .exceptions import ValenceSDKException

logger = get_logger()


class StreamConnection:
    """WebSocket connection for real-time emotion detection."""

    def __init__(self, websocket_url, api_key, model="4emotions"):
        """Initialize a streaming connection.

        Args:
            websocket_url (str): WebSocket URL for streaming
            api_key (str): API key for authentication
            model (str): Model type ('4emotions' or '7emotions')
        """
        self.websocket_url = websocket_url
        self.api_key = api_key
        self.model = model
        self.sio = socketio.Client()
        self._prediction_callback = None
        self._error_callback = None
        self._connected_callback = None
        self._setup_handlers()

    def _setup_handlers(self):
        """Setup Socket.IO event handlers."""

        @self.sio.on('connected')
        def on_connected(data):
            logger.info(f"WebSocket connected: session_id={data.get('session_id')}")
            if self._connected_callback:
                self._connected_callback(data)

        @self.sio.on('prediction')
        def on_prediction(data):
            logger.debug(f"Received prediction: {data.get('main_emotion')}")
            if self._prediction_callback:
                self._prediction_callback(data)

        @self.sio.on('error')
        def on_error(data):
            logger.error(f"WebSocket error: {data.get('error')}")
            if self._error_callback:
                self._error_callback(data)

        @self.sio.on('warning')
        def on_warning(data):
            logger.warning(f"WebSocket warning: {data.get('message')}")

        @self.sio.on('disconnect')
        def on_disconnect():
            logger.info("WebSocket disconnected")

    def connect(self):
        """Establish WebSocket connection.

        Raises:
            ValenceSDKException: If connection fails
        """
        try:
            logger.info(f"Connecting to WebSocket: {self.websocket_url}")
            self.sio.connect(
                self.websocket_url,
                auth={'api_key': self.api_key},
                transports=['polling', 'websocket']
            )
            logger.info("WebSocket connection established")
        except Exception as e:
            logger.error(f"Failed to connect to WebSocket: {e}")
            raise ValenceSDKException(f"WebSocket connection failed: {e}")

    def send_audio(self, audio_data: bytes):
        """Send audio chunk for real-time prediction.

        Args:
            audio_data (bytes): Raw PCM audio data

        Raises:
            ValenceSDKException: If send fails
        """
        if not self.sio.connected:
            raise ValenceSDKException("Not connected to WebSocket")

        try:
            payload = base64.b64encode(audio_data).decode('utf-8')
            self.sio.emit('message', {
                'service': 'emotion',
                'action': 'prediction',
                'model': self.model,
                'payload': payload
            })
            logger.debug(f"Sent audio chunk ({len(audio_data)} bytes)")
        except Exception as e:
            logger.error(f"Failed to send audio: {e}")
            raise ValenceSDKException(f"Failed to send audio: {e}")

    def on_prediction(self, callback):
        """Register callback for prediction events.

        Args:
            callback: Function to call when prediction is received
                     Signature: callback(data: dict)
        """
        self._prediction_callback = callback

    def on_error(self, callback):
        """Register callback for error events.

        Args:
            callback: Function to call when error occurs
                     Signature: callback(data: dict)
        """
        self._error_callback = callback

    def on_connected(self, callback):
        """Register callback for connection events.

        Args:
            callback: Function to call when connected
                     Signature: callback(data: dict)
        """
        self._connected_callback = callback

    def disconnect(self):
        """Close WebSocket connection."""
        if self.sio.connected:
            logger.info("Disconnecting from WebSocket")
            self.sio.disconnect()
        else:
            logger.debug("Already disconnected")

    def wait(self):
        """Wait for events (blocking). Use in simple scripts."""
        try:
            self.sio.wait()
        except KeyboardInterrupt:
            logger.info("Interrupted by user")
            self.disconnect()

    @property
    def connected(self):
        """Check if WebSocket is connected."""
        return self.sio.connected


class StreamingAPI:
    """Client for WebSocket streaming emotion detection."""

    def __init__(self, api_key, websocket_url):
        """Initialize the Streaming API client.

        Args:
            api_key (str): API key for authentication
            websocket_url (str): WebSocket URL for streaming
        """
        self.api_key = api_key
        self.websocket_url = websocket_url

    def connect(self, model: str = "4emotions"):
        """Create and return a new streaming connection.

        Args:
            model (str): Model type ('4emotions' or '7emotions')

        Returns:
            StreamConnection: Active streaming connection

        Raises:
            ValueError: If model is not '4emotions' or '7emotions'

        Example:
            >>> stream = client.streaming.connect(model="4emotions")
            >>> stream.on_prediction(lambda data: print(data))
            >>> stream.connect()
            >>> stream.send_audio(audio_bytes)
            >>> stream.disconnect()
        """
        valid_models = ["4emotions", "7emotions"]
        if model not in valid_models:
            raise ValueError(f"Invalid model '{model}'. Must be one of: {valid_models}")
        return StreamConnection(self.websocket_url, self.api_key, model)
