"""Streaming iterators for NDJSON and SSE responses."""

from __future__ import annotations

import json
from typing import AsyncIterator, Generic, Iterator, Literal, TypeVar

import httpx
from pydantic import BaseModel

from goodmem._generated.models.retrieve_memory_event import RetrieveMemoryEvent

T = TypeVar("T", bound=BaseModel)

Format = Literal["ndjson", "sse"]

_ACCEPT: dict[str, str] = {
    "ndjson": "application/x-ndjson",
    "sse": "text/event-stream",
}


def _validate_format(format: str) -> Format:
    if format not in _ACCEPT:
        raise ValueError(f"format must be one of {list(_ACCEPT)}, got {format!r}")
    return format  # type: ignore[return-value]


def _parse_sse_lines(lines: Iterator[str], model: type[T]) -> Iterator[T]:
    """Yield model instances from a sequence of SSE text lines."""
    event_type: str | None = None
    pending_data: str | None = None
    for line in lines:
        if line.startswith("event:"):
            event_type = line[6:].strip()
        elif line.startswith("data:"):
            pending_data = line[5:].strip()
        elif not line:  # blank line = event boundary
            if event_type != "close" and pending_data:
                yield model.model_validate(json.loads(pending_data))
            event_type = None
            pending_data = None


async def _aparse_sse_lines(
    lines: AsyncIterator[str], model: type[T],
) -> AsyncIterator[T]:
    """Yield model instances from an async sequence of SSE text lines."""
    event_type: str | None = None
    pending_data: str | None = None
    async for line in lines:
        if line.startswith("event:"):
            event_type = line[6:].strip()
        elif line.startswith("data:"):
            pending_data = line[5:].strip()
        elif not line:  # blank line = event boundary
            if event_type != "close" and pending_data:
                yield model.model_validate(json.loads(pending_data))
            event_type = None
            pending_data = None


class Stream(Generic[T]):
    """Iterable stream of pydantic model instances from an NDJSON or SSE response.

    Usage (NDJSON, default)::

        with client.memories.retrieve(message="...", space_ids=["..."]) as stream:
            for event in stream:
                if event.retrieved_item:
                    print(event.retrieved_item)

    Usage (SSE)::

        with client.memories.retrieve(message="...", space_ids=["..."], format="sse") as stream:
            for event in stream:
                if event.abstract_reply:
                    print(event.abstract_reply.text)
    """

    def __init__(
        self,
        http_client: httpx.Client,
        *,
        method: str,
        url: str,
        model: type[T],
        format: Format = "ndjson",
        **request_kwargs,
    ) -> None:
        self._http = http_client
        self._method = method
        self._url = url
        self._model = model
        self._format = format
        self._request_kwargs = request_kwargs
        self._stream_ctx = None
        self._response: httpx.Response | None = None

    def __enter__(self) -> Stream[T]:
        self._stream_ctx = self._http.stream(
            self._method, self._url, **self._request_kwargs
        )
        self._response = self._stream_ctx.__enter__()
        return self

    def __exit__(self, *args) -> None:
        if self._stream_ctx is not None:
            self._stream_ctx.__exit__(*args)

    def __iter__(self) -> Iterator[T]:
        if self._response is None:
            raise RuntimeError("Use as context manager: with stream as s: ...")
        if self._format == "sse":
            yield from _parse_sse_lines(self._response.iter_lines(), self._model)
        else:
            for line in self._response.iter_lines():
                line = line.strip()
                if not line:
                    continue
                yield self._model.model_validate(json.loads(line))


class AsyncStream(Generic[T]):
    """Async iterable stream of pydantic model instances from an NDJSON or SSE response.

    Usage (NDJSON, default)::

        async with client.memories.retrieve(message="...", space_ids=["..."]) as stream:
            async for event in stream:
                if event.retrieved_item:
                    print(event.retrieved_item)

    Usage (SSE)::

        async with client.memories.retrieve(message="...", space_ids=["..."], format="sse") as stream:
            async for event in stream:
                if event.abstract_reply:
                    print(event.abstract_reply.text)
    """

    def __init__(
        self,
        http_client: httpx.AsyncClient,
        *,
        method: str,
        url: str,
        model: type[T],
        format: Format = "ndjson",
        **request_kwargs,
    ) -> None:
        self._http = http_client
        self._method = method
        self._url = url
        self._model = model
        self._format = format
        self._request_kwargs = request_kwargs
        self._stream_ctx = None
        self._response: httpx.Response | None = None

    async def __aenter__(self) -> AsyncStream[T]:
        self._stream_ctx = self._http.stream(
            self._method, self._url, **self._request_kwargs
        )
        self._response = await self._stream_ctx.__aenter__()
        return self

    async def __aexit__(self, *args) -> None:
        if self._stream_ctx is not None:
            await self._stream_ctx.__aexit__(*args)

    async def __aiter__(self) -> AsyncIterator[T]:
        if self._response is None:
            raise RuntimeError("Use as async context manager: async with stream as s: ...")
        if self._format == "sse":
            async for event in _aparse_sse_lines(self._response.aiter_lines(), self._model):
                yield event
        else:
            async for line in self._response.aiter_lines():
                line = line.strip()
                if not line:
                    continue
                yield self._model.model_validate(json.loads(line))


# Backwards-compatible aliases for existing code that references the old names.
RetrieveMemoryStream = Stream[RetrieveMemoryEvent]
AsyncRetrieveMemoryStream = AsyncStream[RetrieveMemoryEvent]
