"""SSE streaming response handlers (sync and async)"""

from __future__ import annotations

import json
from typing import Iterator, TYPE_CHECKING

if TYPE_CHECKING:
    import httpx
    from ..models.chat import StreamEvent


_DONE_SENTINEL = "data: [DONE]"
_DATA_PREFIX = "data: "


def _parse_sse_line(line: str) -> "StreamEvent | None":
    """Parse a single SSE line into a StreamEvent.

    Args:
        line: Raw SSE line from the server

    Returns:
        StreamEvent if the line carries data, None otherwise
    """
    from ..models.chat import StreamEvent

    line = line.strip()
    if not line or line == _DONE_SENTINEL:
        return None
    if not line.startswith(_DATA_PREFIX):
        return None

    data = line[len(_DATA_PREFIX):]
    try:
        payload = json.loads(data)
    except json.JSONDecodeError:
        return None

    return StreamEvent(**payload)


class Stream:
    """Synchronous SSE stream iterator

    Lazily iterates over Server-Sent Events from the chat endpoint.
    The underlying HTTP connection is opened on first iteration and
    closed automatically when iteration completes or an exception occurs.

    Example::

        stream = client.chat.create(model="openai/gpt-5.4", messages=[...], stream=True)
        for event in stream:
            if event.type == "content" and event.choices:
                print(event.choices[0].delta.content or "", end="", flush=True)

    Attributes:
        _http_client: Underlying httpx.Client instance
        _method: HTTP method
        _url: Request URL
        _kwargs: Additional httpx request kwargs
    """

    def __init__(self, http_client: "httpx.Client", method: str, url: str, **kwargs):
        self._http_client = http_client
        self._method = method
        self._url = url
        self._kwargs = kwargs

    def __iter__(self) -> Iterator["StreamEvent"]:
        """Iterate over stream events."""
        from .exceptions import OctenStreamError, OctenAPIError, OctenAuthenticationError, OctenRateLimitError

        with self._http_client.stream(
            self._method, self._url, **self._kwargs
        ) as response:
            # Non-2xx before streaming starts means a hard HTTP error
            if not response.is_success:
                body = response.read().decode("utf-8", errors="replace")
                request_id = response.headers.get("x-request-id")
                try:
                    payload = json.loads(body)
                    msg = payload.get("msg") or payload.get("message") or body
                except Exception:
                    msg = body or f"HTTP {response.status_code}"

                if response.status_code == 401:
                    raise OctenAuthenticationError(message=msg, request_id=request_id)
                if response.status_code == 429:
                    retry_after_hdr = response.headers.get("retry-after")
                    retry_after = int(retry_after_hdr) if retry_after_hdr else None
                    raise OctenRateLimitError(message=msg, retry_after=retry_after, request_id=request_id)
                raise OctenAPIError(
                    message=msg,
                    status_code=response.status_code,
                    request_id=request_id,
                )

            for line in response.iter_lines():
                event = _parse_sse_line(line)
                if event is None:
                    continue
                yield event
                # In-stream error events are yielded first, then raised
                if event.error:
                    raise OctenStreamError(
                        message=event.error.message,
                        code=event.error.code,
                    )


class AsyncStream:
    """Asynchronous SSE stream iterator

    Async version of :class:`Stream`. Use with ``async for``.

    Example::

        stream = await client.chat.create(model="openai/gpt-5.4", messages=[...], stream=True)
        async for event in stream:
            if event.type == "content" and event.choices:
                print(event.choices[0].delta.content or "", end="", flush=True)

    Attributes:
        _http_client: Underlying httpx.AsyncClient instance
        _method: HTTP method
        _url: Request URL
        _kwargs: Additional httpx request kwargs
    """

    def __init__(self, http_client: "httpx.AsyncClient", method: str, url: str, **kwargs):
        self._http_client = http_client
        self._method = method
        self._url = url
        self._kwargs = kwargs

    async def __aiter__(self):
        """Asynchronously iterate over stream events."""
        from .exceptions import OctenStreamError, OctenAPIError, OctenAuthenticationError, OctenRateLimitError

        async with self._http_client.stream(
            self._method, self._url, **self._kwargs
        ) as response:
            if not response.is_success:
                body = (await response.aread()).decode("utf-8", errors="replace")
                request_id = response.headers.get("x-request-id")
                try:
                    payload = json.loads(body)
                    msg = payload.get("msg") or payload.get("message") or body
                except Exception:
                    msg = body or f"HTTP {response.status_code}"

                if response.status_code == 401:
                    raise OctenAuthenticationError(message=msg, request_id=request_id)
                if response.status_code == 429:
                    retry_after_hdr = response.headers.get("retry-after")
                    retry_after = int(retry_after_hdr) if retry_after_hdr else None
                    raise OctenRateLimitError(message=msg, retry_after=retry_after, request_id=request_id)
                raise OctenAPIError(
                    message=msg,
                    status_code=response.status_code,
                    request_id=request_id,
                )

            async for line in response.aiter_lines():
                event = _parse_sse_line(line)
                if event is None:
                    continue
                yield event
                if event.error:
                    raise OctenStreamError(
                        message=event.error.message,
                        code=event.error.code,
                    )
