"""Server-Sent Events (SSE) extractor for real-time data streams.

Uses httpx (already a core dependency) to connect to SSE endpoints
and yield batches of parsed events.
"""

from __future__ import annotations

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

import httpx

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

logger = logging.getLogger(__name__)


class SSEExtractor(BaseExtractor):
    """Extract data from a Server-Sent Events stream.

    Connects to an SSE endpoint and yields batches of parsed events.
    Supports automatic reconnection with Last-Event-ID header.

    Args:
        url: SSE endpoint URL.
        headers: Optional HTTP headers.
        event_filter: Only yield events matching this event type.
        data_format: How to parse event data — 'json' or 'text'.
        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,
        event_filter: str | None = None,
        data_format: str = "json",
        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._event_filter = event_filter
        self._data_format = data_format
        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]) -> SSEExtractor:
        """Create an SSEExtractor from a configuration dictionary.

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

        Returns:
            Configured SSEExtractor instance.
        """
        return cls(
            url=config["url"],
            headers=config.get("headers"),
            event_filter=config.get("event_filter"),
            data_format=config.get("data_format", "json"),
            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 SSE-specific configuration."""
        if not extract_config.get("url"):
            raise ValueError("SSE extractor requires 'url'")

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

        Yields:
            Batches of parsed SSE events.
        """
        cfg = self._streaming_config
        accumulator = BatchAccumulator(cfg.batch_size, cfg.batch_timeout)
        total_records = 0
        total_batches = 0
        last_event_id: str | None = None
        reconnect_count = 0

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

            try:
                async for record, event_id in self._stream_events(last_event_id):
                    if event_id:
                        last_event_id = event_id

                    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

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

            except (
                httpx.RemoteProtocolError,
                httpx.ReadError,
                httpx.ConnectError,
            ) as e:
                reconnect_count += 1
                if not self._reconnect:
                    logger.error("SSE connection failed (reconnect disabled): %s", e)
                    break
                if self._max_reconnects and reconnect_count > self._max_reconnects:
                    logger.error(
                        "SSE max reconnects (%s) exceeded", self._max_reconnects
                    )
                    break
                logger.warning(
                    "SSE 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

    async def _stream_events(
        self, last_event_id: str | None
    ) -> AsyncIterator[tuple[dict[str, Any], str | None]]:
        """Connect to SSE stream and yield parsed events."""
        headers = dict(self._headers)
        headers["Accept"] = "text/event-stream"
        if last_event_id:
            headers["Last-Event-ID"] = last_event_id

        async with httpx.AsyncClient() as client:
            async with client.stream("GET", self._url, headers=headers) as response:
                response.raise_for_status()
                event_data = ""
                event_type = ""
                event_id: str | None = None

                async for line in response.aiter_lines():
                    if line.startswith("data:"):
                        data_value = line[5:].strip()
                        event_data = (
                            event_data + "\n" + data_value if event_data else data_value
                        )
                    elif line.startswith("event:"):
                        event_type = line[6:].strip()
                    elif line.startswith("id:"):
                        event_id = line[3:].strip()
                    elif line.startswith("retry:"):
                        pass  # Server retry hint — not used
                    elif line == "":
                        # Empty line = end of event
                        if event_data:
                            if (
                                self._event_filter is None
                                or event_type == self._event_filter
                            ):
                                record = self._parse_data(
                                    event_data, event_type, event_id
                                )
                                yield record, event_id
                            event_data = ""
                            event_type = ""
                            event_id = None

    def _parse_data(
        self, data: str, event_type: str, event_id: str | None
    ) -> dict[str, Any]:
        """Parse SSE event data into a record dict."""
        if self._data_format == "json":
            try:
                parsed = json.loads(data)
                if isinstance(parsed, dict):
                    if event_type:
                        parsed["_event_type"] = event_type
                    if event_id:
                        parsed["_event_id"] = event_id
                    return parsed
                return {
                    "data": parsed,
                    "_event_type": event_type,
                    "_event_id": event_id,
                }
            except json.JSONDecodeError:
                return {
                    "data": data,
                    "_event_type": event_type,
                    "_event_id": event_id,
                    "_parse_error": True,
                }

        return {"data": data, "_event_type": event_type, "_event_id": event_id}

    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
