"""Streaming iterators for NDJSON responses."""

from __future__ import annotations

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

import httpx
from pydantic import BaseModel

from goodmem.errors import GoodMemError, raise_for_status
from goodmem.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]


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

    Usage::

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

    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__()
        try:
            if self._response.status_code >= 400:
                self._response.read()  # read error body so .text works
            raise_for_status(self._response)
        except Exception:
            self._stream_ctx.__exit__(None, None, None)
            raise
        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 != "ndjson":
            raise NotImplementedError(f"Streaming format {self._format!r} is not yet supported")
        for line in self._response.iter_lines():
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
            except json.JSONDecodeError as e:
                raise GoodMemError(f"Failed to parse streaming response line: {e}") from e
            yield self._model.model_validate(data)


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

    Usage::

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

    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__()
        try:
            if self._response.status_code >= 400:
                await self._response.aread()  # read error body so .text works
            raise_for_status(self._response)
        except Exception:
            await self._stream_ctx.__aexit__(None, None, None)
            raise
        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 != "ndjson":
            raise NotImplementedError(f"Streaming format {self._format!r} is not yet supported")
        async for line in self._response.aiter_lines():
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
            except json.JSONDecodeError as e:
                raise GoodMemError(f"Failed to parse streaming response line: {e}") from e
            yield self._model.model_validate(data)


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