# This file was auto-generated by Fern from our API Definition.

import codecs
import re
import time
from contextlib import asynccontextmanager, contextmanager
from typing import (
    Any,
    AsyncContextManager,
    AsyncGenerator,
    AsyncIterator,
    Callable,
    ContextManager,
    Iterator,
    Optional,
)

import anyio
import httpx
from ._decoders import SSEDecoder
from ._exceptions import SSEError
from ._models import ServerSentEvent

MAX_LINE_SIZE: int = 1_048_576  # 1 MiB

# Reconnection defaults, mirroring the TypeScript SDK's Stream implementation.
DEFAULT_MAX_RECONNECTION_ATTEMPTS: int = 5
DEFAULT_RECONNECT_DELAY_MS: int = 1_000
MAX_RECONNECT_DELAY_MS: int = 30_000


# A reconnect callback re-issues the original request (with a ``Last-Event-ID``
# header set to the supplied event id) and returns a *context manager* yielding
# a fresh streaming ``httpx.Response``. Sync clients supply a sync context
# manager; async clients supply an async one.
class EventSource:
    def __init__(
        self,
        response: httpx.Response,
        *,
        resumable: bool = False,
        stream_reconnection_enabled: bool = True,
        max_stream_reconnection_attempts: Optional[int] = None,
        stream_terminator: Optional[str] = None,
        reconnect: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._response = response
        self._resumable = resumable
        self._stream_reconnection_enabled = stream_reconnection_enabled
        self._max_stream_reconnection_attempts = max_stream_reconnection_attempts
        self._stream_terminator = stream_terminator
        self._reconnect = reconnect

    @staticmethod
    def _is_event_stream(response: httpx.Response) -> bool:
        content_type = response.headers.get("content-type", "").partition(";")[0]
        return "text/event-stream" in content_type

    def _check_content_type(self) -> None:
        if not self._is_event_stream(self._response):
            content_type = self._response.headers.get("content-type", "").partition(";")[0]
            raise SSEError(
                f"Expected response header Content-Type to contain 'text/event-stream', got {content_type!r}"
            )

    def _is_reconnect_response_usable(self, response: httpx.Response) -> bool:
        """Whether a reconnected response can be resumed as an SSE stream.

        ``httpx.stream`` does not raise on non-success status, so a resume that
        returns an error page (e.g. ``200 text/html`` or a ``500`` body) would
        otherwise be parsed as SSE and yield garbage/zero events. Such a
        response is treated as a failed attempt (back off and retry) instead.
        """
        return response.status_code < 400 and self._is_event_stream(response)

    def _get_charset(self, response: Optional[httpx.Response] = None) -> str:
        """Extract charset from Content-Type header, fallback to UTF-8."""
        resolved = response if response is not None else self._response
        content_type = resolved.headers.get("content-type", "")

        # Parse charset parameter using regex
        charset_match = re.search(r"charset=([^;\s]+)", content_type, re.IGNORECASE)
        if charset_match:
            charset = charset_match.group(1).strip("\"'")
            # Validate that it's a known encoding
            try:
                # Test if the charset is valid by trying to encode/decode
                "test".encode(charset).decode(charset)
                return charset
            except (LookupError, UnicodeError):
                # If charset is invalid, fall back to UTF-8
                pass

        # Default to UTF-8 if no charset specified or invalid charset
        return "utf-8"

    @property
    def response(self) -> httpx.Response:
        return self._response

    @staticmethod
    def _normalize_sse_line_endings(buf: str) -> str:
        """Normalize line endings per the SSE spec (\\r\\n → \\n, bare \\r → \\n).

        A trailing \\r is preserved because it may pair with a leading \\n in
        the next chunk to form a single \\r\\n terminator.
        """
        buf = buf.replace("\r\n", "\n")
        if buf.endswith("\r"):
            return buf[:-1].replace("\r", "\n") + "\r"
        return buf.replace("\r", "\n")

    def _new_text_decoder(self, response: Optional[httpx.Response] = None) -> "codecs.IncrementalDecoder":
        return codecs.getincrementaldecoder(self._get_charset(response))(errors="replace")

    def _reconnect_applicable(self) -> bool:
        """Whether reconnection is configured for this stream at all.

        This is the terminator-gating half of the reconnect decision, kept
        separate from :meth:`_should_reconnect` (which additionally requires a
        last *dispatched* id and an unexhausted attempt budget). The split lets
        a mid-stream transport error terminate consistently:
        - a stream that can never reconnect (non-resumable, no terminator,
          disabled, or no callback) must re-raise the error to the caller, so a
          truncated stream is not mistaken for a clean completion;
        - a resumable stream that has merely run out of attempts (or has no id
          to resume from) ends cleanly — the same way an exhausted empty/error
          -body resume already does, matching the TypeScript ``return``.
        """
        return (
            self._resumable
            and self._stream_terminator is not None
            and self._stream_reconnection_enabled
            and self._reconnect is not None
        )

    def _should_reconnect(self, last_dispatched_id: Optional[str], reconnect_attempts: int) -> bool:
        """Decide whether a prematurely-ended stream should be reconnected.

        Mirrors the TypeScript ``shouldReconnect`` gating:
        - only resumable SSE endpoints with a configured terminator, reconnect
          enabled, and a reconnect callback are eligible (see
          :meth:`_reconnect_applicable`);
        - a last *dispatched* event id must exist to resume from;
        - the consecutive-failed-attempt cap must not be exceeded.
        """
        if not self._reconnect_applicable():
            return False
        if not last_dispatched_id:
            return False
        max_attempts = (
            self._max_stream_reconnection_attempts
            if self._max_stream_reconnection_attempts is not None
            else DEFAULT_MAX_RECONNECTION_ATTEMPTS
        )
        if reconnect_attempts >= max_attempts:
            return False
        return True

    def _reconnect_delay_seconds(self, last_retry: Optional[int]) -> float:
        """Backoff before a reconnect.

        Uses the server's most recent ``retry:`` directive (milliseconds) when
        present, otherwise a default of ``DEFAULT_RECONNECT_DELAY_MS``, clamped
        to ``MAX_RECONNECT_DELAY_MS``.
        """
        base_ms = last_retry if (last_retry is not None and last_retry > 0) else DEFAULT_RECONNECT_DELAY_MS
        return min(base_ms, MAX_RECONNECT_DELAY_MS) / 1000.0

    def _sleep_before_reconnect(self, last_retry: Optional[int]) -> None:
        # ``time.sleep`` blocks the calling thread but remains interruptible by
        # signals (e.g. ``KeyboardInterrupt``), which propagate out and abort
        # the reconnect without issuing another request.
        time.sleep(self._reconnect_delay_seconds(last_retry))

    async def _asleep_before_reconnect(self, last_retry: Optional[int]) -> None:
        # ``anyio.sleep`` is cancellation-aware: if the consumer cancels the task
        # or closes the async generator mid-delay, this raises (and no further
        # request is issued) instead of blocking for the whole interval.
        await anyio.sleep(self._reconnect_delay_seconds(last_retry))

    def _decode_response(
        self,
        response: httpx.Response,
        decoder: SSEDecoder,
        text_decoder: "codecs.IncrementalDecoder",
    ) -> Iterator[ServerSentEvent]:
        buf = ""
        for chunk in response.iter_bytes():
            buf += text_decoder.decode(chunk)
            buf = self._normalize_sse_line_endings(buf)

            while "\n" in buf:
                line, buf = buf.split("\n", 1)
                sse = decoder.decode(line)
                if sse is not None:
                    yield sse

            if len(buf) > MAX_LINE_SIZE:
                raise SSEError(
                    f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline"
                )

        yield from self._flush_decoder(buf, decoder, text_decoder)

    async def _adecode_response(
        self,
        response: httpx.Response,
        decoder: SSEDecoder,
        text_decoder: "codecs.IncrementalDecoder",
    ) -> AsyncGenerator[ServerSentEvent, None]:
        buf = ""
        async for chunk in response.aiter_bytes():
            buf += text_decoder.decode(chunk)
            buf = self._normalize_sse_line_endings(buf)

            while "\n" in buf:
                line, buf = buf.split("\n", 1)
                sse = decoder.decode(line)
                if sse is not None:
                    yield sse

            if len(buf) > MAX_LINE_SIZE:
                raise SSEError(
                    f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline"
                )

        for sse in self._flush_decoder(buf, decoder, text_decoder):
            yield sse

    def _flush_decoder(
        self,
        buf: str,
        decoder: SSEDecoder,
        text_decoder: "codecs.IncrementalDecoder",
    ) -> Iterator[ServerSentEvent]:
        # Flush any remaining bytes from the incremental decoder
        buf += text_decoder.decode(b"", final=True)
        buf = buf.replace("\r\n", "\n").replace("\r", "\n")

        if len(buf) > MAX_LINE_SIZE:
            raise SSEError(
                f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline"
            )

        while "\n" in buf:
            line, buf = buf.split("\n", 1)
            sse = decoder.decode(line)
            if sse is not None:
                yield sse

        if buf.strip():
            sse = decoder.decode(buf)
            if sse is not None:
                yield sse

    def iter_sse(self) -> Iterator[ServerSentEvent]:
        self._check_content_type()
        decoder = SSEDecoder()
        text_decoder = self._new_text_decoder()

        last_dispatched_id: Optional[str] = None
        last_retry: Optional[int] = None
        # Consecutive failed reconnection attempts. Reset to 0 whenever an event
        # is successfully dispatched (reset-on-progress) — matching browser
        # `EventSource` semantics: a server that emits >=1 event then drops on
        # every connection can reconnect indefinitely.
        reconnect_attempts = 0

        # ``None`` means there is no live stream to read this iteration (e.g. a
        # failed reconnect); the loop then re-evaluates the reconnect decision
        # without re-reading an exhausted response.
        response: Optional[httpx.Response] = self._response
        # Context manager for a response we opened ourselves and must close.
        # The initial response is owned by the caller, so it starts as None.
        owned_cm: Optional[ContextManager[httpx.Response]] = None
        try:
            while True:
                if response is not None:
                    events = self._decode_response(response, decoder, text_decoder)
                    while True:
                        try:
                            sse = next(events)
                        except StopIteration:
                            break
                        except SSEError:
                            # A protocol violation (e.g. an oversized line) is a
                            # genuine error, not a dropped connection; propagate it.
                            # Listed first because ``SSEError`` subclasses
                            # ``httpx.TransportError``.
                            raise
                        except httpx.TransportError:
                            # A transport error mid-stream (e.g. the server dropped
                            # the connection: ``ReadError``/``RemoteProtocolError``)
                            # is a premature end. Only swallow it when reconnection
                            # is configured for this stream; otherwise re-raise so a
                            # non-resumable stream still surfaces the error to the
                            # caller instead of looking like a clean completion.
                            # When reconnection is applicable but the attempt budget
                            # is exhausted, we ``break`` and end cleanly below — the
                            # same way an exhausted empty/error-body resume does, so
                            # give-up is consistent regardless of failure shape.
                            # ``next`` is used rather than ``for`` so this cannot
                            # swallow a ``GeneratorExit`` raised at a ``yield``.
                            if not self._reconnect_applicable():
                                raise
                            break
                        yield sse
                        if sse.id:
                            last_dispatched_id = sse.id
                        if sse.retry is not None:
                            last_retry = sse.retry
                        reconnect_attempts = 0

                if not self._should_reconnect(last_dispatched_id, reconnect_attempts):
                    return
                reconnect_attempts += 1

                self._sleep_before_reconnect(last_retry)

                # Close the previously-opened reconnect response before opening
                # a new one so we never hold more than one extra connection.
                if owned_cm is not None:
                    owned_cm.__exit__(None, None, None)
                    owned_cm = None

                assert self._reconnect is not None  # guaranteed by _should_reconnect
                try:
                    cm: ContextManager[httpx.Response] = self._reconnect(last_dispatched_id or "")
                    new_response = cm.__enter__()
                except Exception:
                    # A failed reconnect consumes an attempt; back off and retry.
                    response = None
                    continue
                owned_cm = cm
                if new_response is None or not self._is_reconnect_response_usable(new_response):
                    # Null/empty body or a non-SSE/error response (e.g. 204/304,
                    # a 500, or an HTML error page): treat as a failed attempt.
                    response = None
                    continue

                response = new_response
                # Drop any partial event left over from the dropped stream, but
                # keep the last event id (per the SSE spec) and start a fresh
                # incremental text decoder for the new connection.
                decoder.reset_in_progress_event()
                text_decoder = self._new_text_decoder(new_response)
        finally:
            if owned_cm is not None:
                owned_cm.__exit__(None, None, None)

    async def aiter_sse(self) -> AsyncGenerator[ServerSentEvent, None]:
        self._check_content_type()
        decoder = SSEDecoder()
        text_decoder = self._new_text_decoder()

        last_dispatched_id: Optional[str] = None
        last_retry: Optional[int] = None
        reconnect_attempts = 0

        response: Optional[httpx.Response] = self._response
        owned_cm: Optional[AsyncContextManager[httpx.Response]] = None
        try:
            while True:
                if response is not None:
                    events = self._adecode_response(response, decoder, text_decoder)
                    while True:
                        try:
                            sse = await events.__anext__()
                        except StopAsyncIteration:
                            break
                        except SSEError:
                            # A protocol violation (e.g. an oversized line) is a
                            # genuine error, not a dropped connection; propagate it.
                            # Listed first because ``SSEError`` subclasses
                            # ``httpx.TransportError``.
                            raise
                        except httpx.TransportError:
                            # A transport error mid-stream (e.g. the server dropped
                            # the connection: ``ReadError``/``RemoteProtocolError``)
                            # is a premature end. Only swallow it when reconnection
                            # is configured for this stream; otherwise re-raise so a
                            # non-resumable stream still surfaces the error to the
                            # caller instead of looking like a clean completion.
                            # When reconnection is applicable but the attempt budget
                            # is exhausted, we ``break`` and end cleanly below — the
                            # same way an exhausted empty/error-body resume does, so
                            # give-up is consistent regardless of failure shape.
                            if not self._reconnect_applicable():
                                raise
                            break
                        yield sse
                        if sse.id:
                            last_dispatched_id = sse.id
                        if sse.retry is not None:
                            last_retry = sse.retry
                        reconnect_attempts = 0

                if not self._should_reconnect(last_dispatched_id, reconnect_attempts):
                    return
                reconnect_attempts += 1

                await self._asleep_before_reconnect(last_retry)

                if owned_cm is not None:
                    await owned_cm.__aexit__(None, None, None)
                    owned_cm = None

                assert self._reconnect is not None  # guaranteed by _should_reconnect
                try:
                    cm: AsyncContextManager[httpx.Response] = self._reconnect(last_dispatched_id or "")
                    new_response = await cm.__aenter__()
                except Exception:
                    response = None
                    continue
                owned_cm = cm
                if new_response is None or not self._is_reconnect_response_usable(new_response):
                    response = None
                    continue

                response = new_response
                decoder.reset_in_progress_event()
                text_decoder = self._new_text_decoder(new_response)
        finally:
            if owned_cm is not None:
                # Shield the close so a cancellation delivered while reading a
                # reconnected response still fully tears the connection down
                # instead of leaking it until the client is closed.
                with anyio.CancelScope(shield=True):
                    await owned_cm.__aexit__(None, None, None)


@contextmanager
def connect_sse(client: httpx.Client, method: str, url: str, **kwargs: Any) -> Iterator[EventSource]:
    headers = kwargs.pop("headers", {})
    headers["Accept"] = "text/event-stream"
    headers["Cache-Control"] = "no-store"

    with client.stream(method, url, headers=headers, **kwargs) as response:
        yield EventSource(response)


@asynccontextmanager
async def aconnect_sse(
    client: httpx.AsyncClient,
    method: str,
    url: str,
    **kwargs: Any,
) -> AsyncIterator[EventSource]:
    headers = kwargs.pop("headers", {})
    headers["Accept"] = "text/event-stream"
    headers["Cache-Control"] = "no-store"

    async with client.stream(method, url, headers=headers, **kwargs) as response:
        yield EventSource(response)
