"""Streaming iterators for NDJSON memory retrieval responses."""

from __future__ import annotations

import json
from typing import AsyncIterator, Iterator

import httpx

from goodmem._generated.models.retrieve_memory_event import RetrieveMemoryEvent


class RetrieveMemoryStream:
    """Iterable stream of RetrieveMemoryEvent objects from 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,
        **request_kwargs,
    ) -> None:
        self._http = http_client
        self._method = method
        self._url = url
        self._request_kwargs = request_kwargs
        self._stream_ctx = None
        self._response: httpx.Response | None = None

    def __enter__(self) -> RetrieveMemoryStream:
        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[RetrieveMemoryEvent]:
        if self._response is None:
            raise RuntimeError("Use as context manager: with stream as s: ...")
        for line in self._response.iter_lines():
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            yield RetrieveMemoryEvent.model_validate(data)


class AsyncRetrieveMemoryStream:
    """Async iterable stream of RetrieveMemoryEvent objects from 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,
        **request_kwargs,
    ) -> None:
        self._http = http_client
        self._method = method
        self._url = url
        self._request_kwargs = request_kwargs
        self._stream_ctx = None
        self._response: httpx.Response | None = None

    async def __aenter__(self) -> AsyncRetrieveMemoryStream:
        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[RetrieveMemoryEvent]:
        if self._response is None:
            raise RuntimeError("Use as async context manager: async with stream as s: ...")
        async for line in self._response.aiter_lines():
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            yield RetrieveMemoryEvent.model_validate(data)
