"""Auto-paginating iterators for list endpoints."""

from __future__ import annotations

from typing import Any, AsyncIterator, Callable, Generic, Iterator, TypeVar

import httpx

from goodmem.errors import raise_for_status

T = TypeVar("T")


class Paginator(Generic[T]):
    """Auto-paginating iterator over list responses.

    Usage::

        # Iterate all items across pages
        for embedder in client.embedders.list():
            print(embedder.display_name)

        # Cap total items returned
        for embedder in client.embedders.list(max_items=10):
            print(embedder.display_name)

        # Get first page only
        page = client.embedders.list().first_page()
        print(page.items, page.next_token)
    """

    def __init__(
        self,
        fetch_fn: Callable[[str | None], httpx.Response],
        response_model: type,
        items_field: str = "items",
        max_items: int | None = None,
    ) -> None:
        self._fetch_fn = fetch_fn
        self._response_model = response_model
        self._items_field = items_field
        self._max_items = max_items

    def __iter__(self) -> Iterator[T]:
        next_token: str | None = None
        yielded = 0
        while True:
            resp = self._fetch_fn(next_token)
            raise_for_status(resp)
            page = self._response_model.model_validate(resp.json())
            items = getattr(page, self._items_field, []) or []
            for item in items:
                if self._max_items is not None and yielded >= self._max_items:
                    return
                yield item
                yielded += 1
            next_token = getattr(page, "next_token", None)
            if not next_token:
                break

    def first_page(self) -> Any:
        """Fetch only the first page (no auto-pagination)."""
        resp = self._fetch_fn(None)
        raise_for_status(resp)
        return self._response_model.model_validate(resp.json())


class AsyncPaginator(Generic[T]):
    """Async auto-paginating iterator over list responses.

    Usage::

        async for embedder in client.embedders.list():
            print(embedder.display_name)
    """

    def __init__(
        self,
        fetch_fn: Callable[[str | None], Any],  # returns Awaitable[httpx.Response]
        response_model: type,
        items_field: str = "items",
        max_items: int | None = None,
    ) -> None:
        self._fetch_fn = fetch_fn
        self._response_model = response_model
        self._items_field = items_field
        self._max_items = max_items

    async def __aiter__(self) -> AsyncIterator[T]:
        next_token: str | None = None
        yielded = 0
        while True:
            resp = await self._fetch_fn(next_token)
            raise_for_status(resp)
            page = self._response_model.model_validate(resp.json())
            items = getattr(page, self._items_field, []) or []
            for item in items:
                if self._max_items is not None and yielded >= self._max_items:
                    return
                yield item
                yielded += 1
            next_token = getattr(page, "next_token", None)
            if not next_token:
                break

    async def first_page(self) -> Any:
        """Fetch only the first page (no auto-pagination)."""
        resp = await self._fetch_fn(None)
        raise_for_status(resp)
        return self._response_model.model_validate(resp.json())
