"""OpenAI-style pagination 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")  # item type


class Page(Generic[T]):
    """A page of results from a list endpoint.

    Usage::

        # Assignment — get first page
        page = client.spaces.list(page_size=5)
        page.data          # [Space, ...] (up to 5)
        page.next_token    # "abc123" or None

        # For loop — auto-paginate all items
        for space in client.spaces.list():
            print(space.name)

        # Cap total items
        for space in client.spaces.list(max_items=20):
            print(space.name)

        # Page-by-page control
        for page in client.spaces.list(page_size=5).iter_pages():
            print(len(page.data), page.next_token)
    """

    def __init__(
        self,
        *,
        data: list[T],
        next_token: str | None,
        fetch_fn: Callable[[str | None], httpx.Response],
        response_model: type,
        items_field: str = "items",
        max_items: int | None = None,
    ) -> None:
        if max_items is not None:
            self.data = data[:max_items]
        else:
            self.data = data
        self.next_token = next_token
        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]:
        yielded = 0
        # Yield items from first page (already fetched)
        for item in self.data:
            if self._max_items is not None and yielded >= self._max_items:
                return
            yield item
            yielded += 1
        # Fetch subsequent pages
        next_token = self.next_token
        while next_token:
            if self._max_items is not None and yielded >= self._max_items:
                return
            resp = self._fetch_fn(next_token)
            raise_for_status(resp)
            page_obj = self._response_model.model_validate(resp.json())
            items = getattr(page_obj, 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_obj, "next_token", None)

    def iter_pages(self) -> Iterator[Page[T]]:
        """Yield page objects one at a time."""
        yield self
        next_token = self.next_token
        remaining = (self._max_items - len(self.data)) if self._max_items is not None else None
        while next_token:
            if remaining is not None and remaining <= 0:
                return
            resp = self._fetch_fn(next_token)
            raise_for_status(resp)
            page_obj = self._response_model.model_validate(resp.json())
            items = getattr(page_obj, self._items_field, []) or []
            nt = getattr(page_obj, "next_token", None)
            page: Page[T] = Page(
                data=items,
                next_token=nt,
                fetch_fn=self._fetch_fn,
                response_model=self._response_model,
                items_field=self._items_field,
                max_items=remaining,
            )
            yield page
            if remaining is not None:
                remaining -= len(page.data)
            next_token = nt


class AsyncPage(Generic[T]):
    """Async version of Page."""

    def __init__(
        self,
        *,
        data: list[T],
        next_token: str | None,
        fetch_fn: Callable[[str | None], Any],  # returns Awaitable[httpx.Response]
        response_model: type,
        items_field: str = "items",
        max_items: int | None = None,
    ) -> None:
        if max_items is not None:
            self.data = data[:max_items]
        else:
            self.data = data
        self.next_token = next_token
        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]:
        yielded = 0
        for item in self.data:
            if self._max_items is not None and yielded >= self._max_items:
                return
            yield item
            yielded += 1
        next_token = self.next_token
        while next_token:
            if self._max_items is not None and yielded >= self._max_items:
                return
            resp = await self._fetch_fn(next_token)
            raise_for_status(resp)
            page_obj = self._response_model.model_validate(resp.json())
            items = getattr(page_obj, 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_obj, "next_token", None)

    async def iter_pages(self) -> AsyncIterator[AsyncPage[T]]:
        """Yield page objects one at a time."""
        yield self
        next_token = self.next_token
        remaining = (self._max_items - len(self.data)) if self._max_items is not None else None
        while next_token:
            if remaining is not None and remaining <= 0:
                return
            resp = await self._fetch_fn(next_token)
            raise_for_status(resp)
            page_obj = self._response_model.model_validate(resp.json())
            items = getattr(page_obj, self._items_field, []) or []
            nt = getattr(page_obj, "next_token", None)
            page: AsyncPage[T] = AsyncPage(
                data=items,
                next_token=nt,
                fetch_fn=self._fetch_fn,
                response_model=self._response_model,
                items_field=self._items_field,
                max_items=remaining,
            )
            yield page
            if remaining is not None:
                remaining -= len(page.data)
            next_token = nt
