"""
neuroflow.stream -- Real-time data streaming over WebSocket.

Provides an async iterator that connects to a NeuroFlow server's
WebSocket endpoint and yields StreamData objects in real time. The
stream handles connection management, automatic reconnection, and
graceful shutdown.

Usage (async):
    async for data in DataStream(ws_url, device_id):
        print(data.eeg)

Usage (sync, via Session.stream()):
    for data in session.stream(metrics=["focus", "alpha_power"]):
        print(data)
"""

from __future__ import annotations

import asyncio
import json
import time
from typing import AsyncIterator, Iterator, Optional

from neuroflow.models import StreamData, StreamMessageType


class DataStream:
    """
    Async iterator for real-time EEG data from a NeuroFlow WebSocket.

    Connects to the server's ``/ws/stream/{device_id}`` endpoint and
    yields parsed ``StreamData`` objects as they arrive. Supports
    configurable subscription channels, downsampling, and automatic
    reconnection on transient failures.

    Args:
        ws_url: Full WebSocket URL (e.g. "ws://localhost:8000/ws/stream/dev_abc123").
        device_id: Device identifier to stream from.
        channels: List of stream message types to subscribe to.
        pipeline_name: Optional pipeline for server-side feature extraction.
        downsample_factor: Only yield every Nth data window (default 1 = all).
        window_seconds: Size of each data window in seconds (default 0.1).
        api_key: Optional API key for authenticated connections.
        max_reconnect_attempts: Number of reconnection attempts before giving up.
        reconnect_delay: Base delay in seconds between reconnection attempts.

    Example::

        stream = DataStream(
            ws_url="ws://localhost:8000/ws/stream/dev_abc123",
            device_id="dev_abc123",
            channels=["eeg_data", "features"],
            pipeline_name="meditation",
        )
        async for data in stream:
            print(f"Alpha power: {data.features.get('avg_alpha_power', 0)}")
    """

    def __init__(
        self,
        ws_url: str,
        device_id: str,
        channels: Optional[list[str]] = None,
        pipeline_name: str = "",
        downsample_factor: int = 1,
        window_seconds: float = 0.1,
        api_key: Optional[str] = None,
        max_reconnect_attempts: int = 5,
        reconnect_delay: float = 1.0,
    ) -> None:
        self._ws_url = ws_url
        self._device_id = device_id
        self._channels = channels or ["eeg_data"]
        self._pipeline_name = pipeline_name
        self._downsample_factor = downsample_factor
        self._window_seconds = window_seconds
        self._api_key = api_key
        self._max_reconnect_attempts = max_reconnect_attempts
        self._reconnect_delay = reconnect_delay
        self._running = False
        self._websocket = None

    async def _connect(self):
        """Establish WebSocket connection and send subscription."""
        try:
            import websockets
        except ImportError:
            raise ImportError(
                "The 'websockets' package is required for streaming. "
                "Install it with: pip install neuroflow"
            )

        extra_headers = {}
        if self._api_key:
            extra_headers["X-API-Key"] = self._api_key

        self._websocket = await websockets.connect(
            self._ws_url,
            additional_headers=extra_headers if extra_headers else None,
        )

        # Send subscription message
        subscription = {
            "type": "subscribe",
            "data": {
                "device_id": self._device_id,
                "channels": self._channels,
                "pipeline_name": self._pipeline_name,
                "downsample_factor": self._downsample_factor,
                "window_seconds": self._window_seconds,
            },
        }
        await self._websocket.send(json.dumps(subscription))

    async def _disconnect(self) -> None:
        """Close the WebSocket connection gracefully."""
        if self._websocket is not None:
            try:
                await self._websocket.close()
            except Exception:
                pass
            self._websocket = None

    def _parse_message(self, raw_message: str) -> Optional[StreamData]:
        """Parse a raw WebSocket message into a StreamData object."""
        try:
            msg = json.loads(raw_message)
        except json.JSONDecodeError:
            return None

        msg_type_str = msg.get("type", "")
        try:
            msg_type = StreamMessageType(msg_type_str)
        except ValueError:
            # Unknown message type -- skip
            return None

        # Filter out control messages (ping/pong) -- not user data
        if msg_type in (StreamMessageType.PING, StreamMessageType.PONG):
            return None

        data = msg.get("data", {})
        timestamp = msg.get("timestamp", time.time())

        stream_data = StreamData(
            timestamp=timestamp,
            message_type=msg_type,
            metadata=data,
        )

        if msg_type == StreamMessageType.EEG_DATA:
            stream_data.eeg = data.get("eeg", [])
            stream_data.channel_names = data.get("channels", [])
            stream_data.num_samples = data.get("samples", 0)
            stream_data.sampling_rate = data.get("sampling_rate", 0)

        elif msg_type == StreamMessageType.FEATURES:
            stream_data.features = data

        elif msg_type == StreamMessageType.QUALITY:
            channel_scores = data.get("channel_scores", [])
            stream_data.overall_quality = data.get("overall", 0.0)
            # Build quality dict from channel score list
            stream_data.quality = {
                f"ch{i}": score for i, score in enumerate(channel_scores)
            }

        elif msg_type == StreamMessageType.PROCESSED_DATA:
            stream_data.eeg = data.get("data", [])

        elif msg_type == StreamMessageType.ERROR:
            stream_data.metadata = data

        return stream_data

    async def __aiter__(self) -> AsyncIterator[StreamData]:
        """
        Async iterate over streaming data.

        Yields StreamData objects as they arrive from the WebSocket.
        Automatically reconnects on transient failures up to
        ``max_reconnect_attempts`` times.
        """
        self._running = True
        attempts = 0

        while self._running:
            try:
                await self._connect()
                attempts = 0  # Reset on successful connection

                async for raw_message in self._websocket:
                    if not self._running:
                        break

                    parsed = self._parse_message(raw_message)
                    if parsed is not None:
                        yield parsed

            except Exception as exc:
                await self._disconnect()

                if not self._running:
                    break

                attempts += 1
                if attempts > self._max_reconnect_attempts:
                    raise ConnectionError(
                        f"Failed to maintain WebSocket connection after "
                        f"{self._max_reconnect_attempts} attempts: {exc}"
                    ) from exc

                delay = self._reconnect_delay * (2 ** (attempts - 1))
                await asyncio.sleep(delay)

            finally:
                if not self._running:
                    await self._disconnect()

    def stop(self) -> None:
        """Signal the stream to stop after the current message."""
        self._running = False

    @property
    def is_running(self) -> bool:
        """Whether the stream is currently active."""
        return self._running


class SyncDataStream:
    """
    Synchronous wrapper around DataStream for use in non-async code.

    Runs the async DataStream in a background thread with its own
    event loop and yields data via a thread-safe queue.

    Args:
        ws_url: Full WebSocket URL.
        device_id: Device identifier.
        channels: Stream message types to subscribe to.
        pipeline_name: Optional pipeline for feature extraction.
        downsample_factor: Yield every Nth data window.
        window_seconds: Data window size in seconds.
        api_key: Optional API key.
        max_items: Maximum items to buffer in the queue.

    Example::

        for data in SyncDataStream("ws://localhost:8000/ws/stream/dev_abc", "dev_abc"):
            print(data.eeg)
    """

    def __init__(
        self,
        ws_url: str,
        device_id: str,
        channels: Optional[list[str]] = None,
        pipeline_name: str = "",
        downsample_factor: int = 1,
        window_seconds: float = 0.1,
        api_key: Optional[str] = None,
        max_items: int = 1000,
    ) -> None:
        import queue
        import threading

        self._async_stream = DataStream(
            ws_url=ws_url,
            device_id=device_id,
            channels=channels,
            pipeline_name=pipeline_name,
            downsample_factor=downsample_factor,
            window_seconds=window_seconds,
            api_key=api_key,
        )
        self._queue: queue.Queue[Optional[StreamData]] = queue.Queue(maxsize=max_items)
        self._thread: Optional[threading.Thread] = None
        self._error: Optional[Exception] = None
        self._stopped = threading.Event()

    def _run_async_loop(self) -> None:
        """Run the async stream in a dedicated event loop."""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        async def _consume() -> None:
            try:
                async for data in self._async_stream:
                    if self._stopped.is_set():
                        break
                    try:
                        self._queue.put_nowait(data)
                    except Exception:
                        # Queue full -- drop oldest and retry
                        try:
                            self._queue.get_nowait()
                            self._queue.put_nowait(data)
                        except Exception:
                            pass
            except Exception as exc:
                self._error = exc
            finally:
                self._queue.put(None)  # Sentinel to signal end

        try:
            loop.run_until_complete(_consume())
        finally:
            loop.close()

    def __iter__(self) -> Iterator[StreamData]:
        """
        Synchronous iterator yielding StreamData objects.

        Starts a background thread with an async event loop to receive
        WebSocket data, then yields items from a thread-safe queue.
        """
        import threading

        self._stopped.clear()
        self._error = None
        self._thread = threading.Thread(
            target=self._run_async_loop,
            name="neuroflow-sync-stream",
            daemon=True,
        )
        self._thread.start()

        try:
            while True:
                item = self._queue.get()
                if item is None:
                    break
                yield item
        finally:
            self.stop()

        if self._error is not None:
            raise self._error

    def stop(self) -> None:
        """Stop the synchronous stream."""
        self._stopped.set()
        self._async_stream.stop()
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=5.0)
            self._thread = None
