"""WebSocket extractor for real-time data streams.

Uses the ``websockets`` library (optional dependency) to connect to
WebSocket endpoints and yield batches of parsed messages.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator

from pycharter.etl_generator.extractors._streaming import (
    BatchAccumulator,
    StreamingConfig,
)
from pycharter.etl_generator.extractors.base import BaseExtractor

logger = logging.getLogger(__name__)


class WebSocketExtractor(BaseExtractor):
    """Extract data from a WebSocket stream.

    Connects to a WebSocket endpoint, optionally sends subscribe messages,
    and yields batches of parsed messages. Supports automatic reconnection.

    Args:
        url: WebSocket endpoint URL (ws:// or wss://).
        headers: Optional connection headers.
        subscribe_messages: Messages to send after connecting (e.g. channel subscriptions).
        data_format: How to parse messages — 'json' or 'text'.
        ping_interval: Seconds between keepalive pings (None to disable).
        batch_size: Records per batch.
        batch_timeout: Seconds before yielding a partial batch.
        reconnect: Whether to reconnect on connection loss.
        reconnect_delay: Seconds to wait before reconnecting.
        max_reconnects: Maximum reconnection attempts (0 = unlimited).
        shutdown_event: External event to signal graceful shutdown.
        max_records: Stop after this many total records.
        max_batches: Stop after yielding this many batches.
    """

    def __init__(
        self,
        url: str,
        headers: dict[str, str] | None = None,
        subscribe_messages: list[dict[str, Any] | str] | None = None,
        data_format: str = "json",
        ping_interval: float | None = 20.0,
        batch_size: int = 1000,
        batch_timeout: float = 5.0,
        reconnect: bool = True,
        reconnect_delay: float = 3.0,
        max_reconnects: int = 10,
        shutdown_event: asyncio.Event | None = None,
        max_records: int | None = None,
        max_batches: int | None = None,
    ) -> None:
        self._url = url
        self._headers = headers or {}
        self._subscribe_messages = subscribe_messages or []
        self._data_format = data_format
        self._ping_interval = ping_interval
        self._streaming_config = StreamingConfig(
            shutdown_event=shutdown_event,
            max_batches=max_batches,
            max_records=max_records,
            batch_timeout=batch_timeout,
            batch_size=batch_size,
        )
        self._reconnect = reconnect
        self._reconnect_delay = reconnect_delay
        self._max_reconnects = max_reconnects

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> WebSocketExtractor:
        """Create a WebSocketExtractor from a configuration dictionary.

        Args:
            config: Extractor configuration with WebSocket-specific keys.

        Returns:
            Configured WebSocketExtractor instance.
        """
        return cls(
            url=config["url"],
            headers=config.get("headers"),
            subscribe_messages=config.get("subscribe_messages"),
            data_format=config.get("data_format", "json"),
            ping_interval=config.get("ping_interval", 20.0),
            batch_size=config.get("batch_size", 1000),
            batch_timeout=config.get("batch_timeout", 5.0),
            reconnect=config.get("reconnect", True),
            reconnect_delay=config.get("reconnect_delay", 3.0),
            max_reconnects=config.get("max_reconnects", 10),
            max_records=config.get("max_records"),
            max_batches=config.get("max_batches"),
        )

    def validate_config(self, extract_config: dict[str, Any]) -> None:
        """Validate WebSocket-specific configuration."""
        if not extract_config.get("url"):
            raise ValueError("WebSocket extractor requires 'url'")

    async def extract(self, **params: Any) -> AsyncIterator[list[dict[str, Any]]]:
        """Extract data from a WebSocket stream.

        Yields:
            Batches of parsed WebSocket messages.

        Raises:
            ImportError: If websockets library is not installed.
        """
        try:
            import websockets  # noqa: F811
        except ImportError:
            raise ImportError(
                "WebSocket extractor requires the 'websockets' package. "
                "Install with: pip install pycharter[streaming] or pip install websockets"
            )

        cfg = self._streaming_config
        accumulator = BatchAccumulator(cfg.batch_size, cfg.batch_timeout)
        total_records = 0
        total_batches = 0
        reconnect_count = 0

        while True:
            if cfg.shutdown_event and cfg.shutdown_event.is_set():
                break

            try:
                connect_kwargs: dict[str, Any] = {
                    "additional_headers": self._headers or None,
                }
                if self._ping_interval is not None:
                    connect_kwargs["ping_interval"] = self._ping_interval

                async with websockets.connect(self._url, **connect_kwargs) as ws:
                    # Send subscribe messages after connecting
                    for msg in self._subscribe_messages:
                        payload = json.dumps(msg) if isinstance(msg, dict) else str(msg)
                        await ws.send(payload)

                    reconnect_count = 0  # Reset on successful connection

                    async for raw_message in ws:
                        record = self._parse_message(raw_message)
                        batch = await accumulator.add(record)
                        total_records += 1

                        if batch:
                            yield batch
                            total_batches += 1

                        if cfg.max_records and total_records >= cfg.max_records:
                            break
                        if cfg.max_batches and total_batches >= cfg.max_batches:
                            break
                        if cfg.shutdown_event and cfg.shutdown_event.is_set():
                            break

                # Connection closed normally — flush and exit
                remaining = await accumulator.flush()
                if remaining:
                    yield remaining
                break

            except Exception as e:
                reconnect_count += 1
                if not self._reconnect:
                    logger.error(
                        "WebSocket connection failed (reconnect disabled): %s", e
                    )
                    break
                if self._max_reconnects and reconnect_count > self._max_reconnects:
                    logger.error(
                        "WebSocket max reconnects (%s) exceeded",
                        self._max_reconnects,
                    )
                    break
                logger.warning(
                    "WebSocket connection lost, reconnecting in %ss (attempt %s): %s",
                    self._reconnect_delay,
                    reconnect_count,
                    e,
                )
                await asyncio.sleep(self._reconnect_delay)

            # Check limits after reconnect
            if cfg.max_records and total_records >= cfg.max_records:
                break
            if cfg.max_batches and total_batches >= cfg.max_batches:
                break

        remaining = await accumulator.flush()
        if remaining:
            yield remaining

    def _parse_message(self, raw_message: str | bytes) -> dict[str, Any]:
        """Parse a WebSocket message into a record dict."""
        if isinstance(raw_message, bytes):
            raw_message = raw_message.decode("utf-8", errors="replace")

        if self._data_format == "json":
            try:
                parsed = json.loads(raw_message)
                if isinstance(parsed, dict):
                    return parsed
                return {"data": parsed}
            except json.JSONDecodeError:
                return {"data": raw_message, "_parse_error": True}

        return {"data": raw_message}

    async def extract_streaming(
        self,
        extract_config: dict[str, Any],
        params: dict[str, Any],
        headers: dict[str, Any],
        contract_dir: Any | None = None,
        batch_size: int = 1000,
        max_records: int | None = None,
        config_context: dict[str, Any] | None = None,
    ) -> AsyncIterator[list[dict[str, Any]]]:
        """Extract via config-driven path (delegates to extract)."""
        async for batch in self.extract(**params):
            yield batch
