"""Streaming support for Celeste."""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Callable, Iterator
from contextlib import suppress
from types import TracebackType
from typing import Any, ClassVar, Self, Unpack

import httpx
from anyio.from_thread import start_blocking_portal

from celeste.exceptions import StreamEventError, StreamNotExhaustedError
from celeste.io import Chunk as ChunkBase
from celeste.io import FinishReason, Output, Usage
from celeste.parameters import Parameters
from celeste.tools import ToolCall
from celeste.types import RawUsage


async def enrich_stream_errors(
    iterator: AsyncIterator[dict[str, Any]],
    error_handler: Callable[[httpx.Response], None],
) -> AsyncIterator[dict[str, Any]]:
    """Wrap stream iterator to enrich HTTP errors with provider-specific messages."""
    try:
        async for event in iterator:
            yield event
    except httpx.HTTPStatusError as e:
        error_handler(e.response)
        raise  # Unreachable — error_handler always raises for error responses


class Stream[Out: Output, Params: Parameters, Chunk: ChunkBase](ABC):
    """Async iterator wrapper providing final Output access after stream exhaustion.

    Supports both async iteration (`async for chunk in stream`) and sync iteration
    (`for chunk in stream`). Sync iteration uses anyio's blocking portal to maintain
    a persistent event loop in a dedicated thread.

    Note: For high-volume scenarios, async iteration is recommended. Sync iteration
    creates a background thread per stream.
    """

    _usage_class: ClassVar[type[Usage]] = Usage
    _finish_reason_class: ClassVar[type[FinishReason]] = FinishReason
    _chunk_class: ClassVar[type[ChunkBase]]
    _output_class: ClassVar[type[Output]]
    _empty_content: ClassVar[Any]
    _error_type_fields: ClassVar[tuple[str, ...]] = ("type", "code")

    def __init__(
        self,
        sse_iterator: AsyncIterator[dict[str, Any]],
        transform_output: Callable[..., Any] | None = None,
        stream_metadata: dict[str, Any] | None = None,
        **parameters: Unpack[Params],  # type: ignore[misc]
    ) -> None:
        """Initialize stream with SSE iterator."""
        self._sse_iterator = sse_iterator
        self._chunks: list[Chunk] = []
        self._closed = False
        self._output: Out | None = None
        self._parameters = parameters
        self._transform_output = transform_output
        self._stream_metadata = stream_metadata or {}
        # Sync iteration state (portal lifecycle managed by __iter__ generator)
        self._sync_generator: Iterator[Chunk] | None = None

    def _build_error_from_value(self, error: Any) -> dict[str, Any]:  # noqa: ANN401
        """Extract {type, message} from an error value using _error_type_fields."""
        if isinstance(error, dict):
            error_type = None
            for field in self._error_type_fields:
                val = error.get(field)
                if val is not None:
                    error_type = str(val)
                    break
            return {
                "type": error_type,
                "message": error.get("message", "Unknown error"),
            }
        # Non-dict error value (e.g., plain string from Cohere)
        return {"message": str(error) if error else "Unknown error"}

    def _parse_stream_error(self, event_data: dict[str, Any]) -> dict[str, Any] | None:
        """Detect error events in the SSE stream.

        Handles two generic SSE error patterns:
        1. Type-based: {"type": "error", "error": {"type": "...", "message": "..."}}
        2. Field-based: {"error": {"message": "...", "type": "..."}}

        Override in provider mixin for non-standard error shapes.
        """
        error = None

        # Pattern 1: Type-based — event has "type": "error" with nested error
        if event_data.get("type") == "error":
            error = event_data.get("error", {})
        # Pattern 2: Field-based — event has top-level "error" dict
        elif isinstance(event_data.get("error"), dict):
            error = event_data["error"]

        if error is None:
            return None

        return self._build_error_from_value(error)

    def _parse_chunk_content(self, event_data: dict[str, Any]) -> Any | None:  # noqa: ANN401
        """Parse content from chunk event. Override in provider mixin."""
        return None

    def _wrap_chunk_content(self, raw_content: Any) -> Any:  # noqa: ANN401
        """Wrap raw content into chunk content type. Override for type transformation."""
        return raw_content

    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
        """Parse SSE event into Chunk (returns None to filter lifecycle events)."""
        error = self._parse_stream_error(event)
        if error is not None:
            raise StreamEventError(
                message=error.get("message", "Unknown stream error"),
                error_type=error.get("type"),
                event_data=event,
                provider=self._stream_metadata.get("provider"),
            )
        content = self._parse_chunk_content(event)
        usage = self._get_chunk_usage(event)
        finish_reason = self._get_chunk_finish_reason(event)
        if content is None:
            if usage is None and finish_reason is None:
                return None
            content = self._empty_content
        else:
            content = self._wrap_chunk_content(content)
        return self._chunk_class(  # type: ignore[return-value]
            content=content,
            finish_reason=finish_reason,
            usage=usage,
            metadata={"event_data": event},
        )

    @abstractmethod
    def _aggregate_content(self, chunks: list[Chunk]) -> Any:  # noqa: ANN401
        """Aggregate content from chunks into raw content (modality-specific)."""
        ...

    def _parse_output(self, chunks: list[Chunk], **parameters: Unpack[Params]) -> Out:  # type: ignore[misc]
        """Parse final Output from accumulated chunks."""
        raw_content = self._aggregate_content(chunks)
        content = (
            self._transform_output(raw_content, **parameters)
            if self._transform_output
            else raw_content
        )
        raw_events = self._aggregate_event_data(chunks)
        return self._output_class(  # type: ignore[return-value]
            content=content,
            usage=self._aggregate_usage(chunks),
            finish_reason=self._aggregate_finish_reason(chunks),
            metadata=self._build_stream_metadata(raw_events),
            tool_calls=self._aggregate_tool_calls(chunks, raw_events),
        )

    def _aggregate_tool_calls(
        self, chunks: list[Chunk], raw_events: list[dict[str, Any]]
    ) -> list[ToolCall]:
        """Aggregate tool calls from stream events. Override in providers that support tools."""
        return []

    def _parse_chunk_usage(self, event_data: dict[str, Any]) -> RawUsage | None:
        """Parse usage from chunk event. Override in provider mixin."""
        return None

    def _parse_chunk_finish_reason(
        self, event_data: dict[str, Any]
    ) -> FinishReason | None:
        """Parse finish reason from chunk event. Override in provider mixin."""
        return None

    def _get_chunk_usage(self, event_data: dict[str, Any]) -> Usage | None:
        """Get modality-typed usage from chunk event."""
        raw = self._parse_chunk_usage(event_data)
        if raw is None:
            return None
        return self._usage_class(**raw)

    def _get_chunk_finish_reason(
        self, event_data: dict[str, Any]
    ) -> FinishReason | None:
        """Get modality-typed finish reason from chunk event."""
        raw = self._parse_chunk_finish_reason(event_data)
        if raw is None:
            return None
        if isinstance(raw, self._finish_reason_class):
            return raw
        return self._finish_reason_class(reason=raw.reason)

    def _build_stream_metadata(
        self, raw_events: list[dict[str, Any]]
    ) -> dict[str, Any]:
        """Build metadata for streaming. Providers override to filter content."""
        return {**self._stream_metadata, "raw_events": raw_events}

    def _aggregate_usage(self, chunks: list[Chunk]) -> Usage:
        """Aggregate usage across chunks (last chunk with usage wins)."""
        for chunk in reversed(chunks):
            if chunk.usage:
                return chunk.usage
        return self._usage_class()

    def _aggregate_finish_reason(self, chunks: list[Chunk]) -> FinishReason | None:
        """Aggregate finish reason across chunks (last chunk with finish_reason wins)."""
        for chunk in reversed(chunks):
            if chunk.finish_reason:
                return chunk.finish_reason
        return None

    def _aggregate_event_data(self, chunks: list[Chunk]) -> list[dict[str, Any]]:
        """Collect raw event_data from chunk metadata."""
        events: list[dict[str, Any]] = []
        for chunk in chunks:
            event_data = chunk.metadata.get("event_data")
            if isinstance(event_data, dict):
                events.append(event_data)
        return events

    def __repr__(self) -> str:
        """Developer-friendly representation showing stream state."""
        if self._output:
            state = "done"
        elif self._closed:
            state = "closed"
        elif self._chunks:
            state = "streaming"
        else:
            state = "idle"

        chunks = f", {len(self._chunks)} chunks" if self._chunks else ""
        return f"<{self.__class__.__name__}: {state}{chunks}>"

    # AsyncIterator protocol
    def __aiter__(self) -> Self:
        """Return self as async iterator."""
        return self

    async def __anext__(self) -> Chunk:
        """Yield next chunk from stream."""
        if self._closed:
            raise StopAsyncIteration

        try:
            async for event in self._sse_iterator:
                chunk = self._parse_chunk(event)
                if chunk is not None:
                    self._chunks.append(chunk)
                    return chunk

            # Stream exhausted naturally
            if self._chunks:
                self._output = self._parse_output(self._chunks, **self._parameters)
            self._closed = True
        except Exception:
            await self.aclose()
            raise

        raise StopAsyncIteration

    # Iterator protocol (sync)
    def __iter__(self) -> Iterator[Chunk]:
        """Sync iterator using a generator with try/finally for guaranteed cleanup.

        Creates a blocking portal in a dedicated thread. The generator's finally
        block ensures the portal is cleaned up on exhaustion, break, exception,
        or garbage collection — unlike __next__ which only cleans up on exhaustion.
        """
        portal_cm = start_blocking_portal()
        portal = portal_cm.__enter__()
        try:
            while True:
                try:
                    yield portal.call(self.__anext__)
                except StopAsyncIteration:
                    return
        finally:
            if not self._closed:
                with suppress(RuntimeError):
                    portal.call(self.aclose)
            portal_cm.__exit__(None, None, None)

    # AsyncContextManager protocol
    async def __aenter__(self) -> Self:
        """Enter context - return self for iteration."""
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> bool:
        """Exit context - ensure cleanup even on exception."""
        await self.aclose()
        return False  # Propagate exceptions

    # ContextManager protocol (sync)
    def __enter__(self) -> Self:
        """Enter sync context - return self for iteration."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Exit sync context. Portal cleanup is handled by __iter__ generator's finally."""
        return

    @property
    def output(self) -> Out:
        """Access final Output after stream exhaustion (raises StreamNotExhaustedError if not ready)."""
        if self._output is None:
            raise StreamNotExhaustedError()
        return self._output

    async def aclose(self) -> None:
        """Explicitly close stream and cleanup resources."""
        if self._closed:
            return

        self._closed = True

        # Fast path: skip if iterator is currently running
        if getattr(self._sse_iterator, "ag_running", False):
            return

        # Close SSE iterator (httpx-sse connection)
        # Use suppress to handle TOCTOU race between ag_running check and aclose
        if hasattr(self._sse_iterator, "aclose"):
            with suppress(RuntimeError):
                await self._sse_iterator.aclose()


__all__ = ["Stream", "enrich_stream_errors"]
