"""Async pagination utilities for automatic cursor and offset-based pagination.

This module provides async iterators that automatically handle pagination,
allowing clients to iterate over all items without managing cursors or offsets.

Example usage:

    # Iterate item-by-item (memory efficient)
    async for item in client.iter_content_parts(agent_id=some_id):
        process(item)

    # Collect all items (default 10k safety limit)
    all_items = await client.iter_content_parts(agent_id=some_id).collect()

    # Collect specific number of items
    items = await client.iter_content_parts().collect(max_items=50)

    # Collect ALL items (no limit)
    all_items = await client.iter_content_parts().collect(max_items=None)

    # Collect with strict limit (raises error if exceeded)
    items = await client.iter_content_parts().collect_strict(max_items=100)

    # Iterate page-by-page
    async for page in client.iter_content_parts().pages():
        process_batch(page)
"""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Awaitable, Callable
from typing import Generic, Self, TypeVar

from pydantic import BaseModel

# Type variables for generic pagination
TItem = TypeVar("TItem")
TResponse = TypeVar("TResponse", bound=BaseModel)

# Default safety limit for collect() to prevent accidental memory issues
DEFAULT_MAX_ITEMS = 10_000


class PaginationLimitExceeded(Exception):
    """Raised when collect() would exceed the safety limit."""

    def __init__(self, limit: int, message: str | None = None):
        self.limit = limit
        super().__init__(
            message
            or f"Pagination limit of {limit} items exceeded. "
            f"Use collect(max_items=None) to disable the limit, "
            f"or specify a higher limit with collect(max_items=N)."
        )


class AsyncPaginatorBase(Generic[TItem], ABC):
    """Base class providing common collection methods for async paginators."""

    _exhausted: bool

    @abstractmethod
    def __aiter__(self) -> Self: ...

    @abstractmethod
    async def __anext__(self) -> TItem: ...

    async def collect(self, max_items: int | None = DEFAULT_MAX_ITEMS) -> list[TItem]:
        """Collect items into a list, stopping when max_items is reached.

        Args:
            max_items: Maximum number of items to collect. Defaults to 10,000 as a
                      safety limit. Pass None to disable the limit entirely and
                      collect all items.

        Returns:
            List of items (up to max_items if specified)
        """
        items: list[TItem] = []
        async for item in self:
            items.append(item)
            if max_items is not None and len(items) >= max_items:
                break
        return items

    async def collect_strict(self, max_items: int = DEFAULT_MAX_ITEMS) -> list[TItem]:
        """Collect all items, raising an error if max_items is exceeded.

        Use this when you want to ensure you're not accidentally fetching
        more data than expected.

        Args:
            max_items: Maximum allowed items before raising an error.

        Returns:
            List of all items from all pages

        Raises:
            PaginationLimitExceeded: If the total items would exceed max_items
        """
        items: list[TItem] = []
        async for item in self:
            items.append(item)
            if len(items) > max_items:
                raise PaginationLimitExceeded(max_items)
        return items


class AsyncCursorPaginator(AsyncPaginatorBase[TItem], Generic[TItem, TResponse]):
    """Async iterator for cursor-based pagination.

    Automatically fetches pages as needed and yields items one at a time.
    Supports collecting all items into a list with optional safety limits.

    Type Parameters:
        TItem: The type of individual items being paginated
        TResponse: The response model type containing the items and cursor
    """

    def __init__(
        self,
        fetch_page: Callable[[str | None], Awaitable[TResponse]],
        extract_items: Callable[[TResponse], list[TItem]],
        extract_cursor: Callable[[TResponse], str | None],
        page_size: int = 500,
    ):
        """Initialize the cursor paginator.

        Args:
            fetch_page: Async function that fetches a page given a cursor (None for first page)
            extract_items: Function to extract items list from response
            extract_cursor: Function to extract next cursor from response (None if no more pages)
            page_size: Number of items to fetch per page (default 500)
        """
        self._fetch_page = fetch_page
        self._extract_items = extract_items
        self._extract_cursor = extract_cursor
        self._page_size = page_size
        self._cursor: str | None = None
        self._buffer: list[TItem] = []
        self._buffer_index: int = 0
        self._exhausted: bool = False

    def __aiter__(self) -> Self:
        """Return self as async iterator."""
        return self

    async def __anext__(self) -> TItem:
        """Get the next item, fetching a new page if needed."""
        # If we have items in buffer, return next one
        if self._buffer_index < len(self._buffer):
            item = self._buffer[self._buffer_index]
            self._buffer_index += 1
            return item

        # Buffer exhausted - try to fetch next page
        if self._exhausted:
            raise StopAsyncIteration

        # Fetch next page
        response = await self._fetch_page(self._cursor)
        self._buffer = self._extract_items(response)
        self._buffer_index = 0
        self._cursor = self._extract_cursor(response)

        # Check if we got any items
        if not self._buffer:
            self._exhausted = True
            raise StopAsyncIteration

        # If no next cursor, this is the last page
        if self._cursor is None:
            self._exhausted = True

        # Return first item from new buffer
        item = self._buffer[self._buffer_index]
        self._buffer_index += 1
        return item

    async def pages(self) -> AsyncIterator[list[TItem]]:
        """Iterate page-by-page instead of item-by-item.

        Yields:
            List of items for each page
        """
        while not self._exhausted:
            response = await self._fetch_page(self._cursor)
            items = self._extract_items(response)
            self._cursor = self._extract_cursor(response)

            if not items:
                self._exhausted = True
                return

            if self._cursor is None:
                self._exhausted = True

            yield items


class AsyncOffsetPaginator(AsyncPaginatorBase[TItem], Generic[TItem, TResponse]):
    """Async iterator for offset-based pagination.

    Automatically fetches pages as needed and yields items one at a time.
    Supports collecting all items into a list with optional safety limits.

    Type Parameters:
        TItem: The type of individual items being paginated
        TResponse: The response model type containing the items and total count
    """

    def __init__(
        self,
        fetch_page: Callable[[int, int], Awaitable[TResponse]],
        extract_items: Callable[[TResponse], list[TItem]],
        extract_total: Callable[[TResponse], int],
        page_size: int = 500,
    ):
        """Initialize the offset paginator.

        Args:
            fetch_page: Async function that fetches a page given (limit, offset)
            extract_items: Function to extract items list from response
            extract_total: Function to extract total count from response
            page_size: Number of items to fetch per page (default 500)
        """
        self._fetch_page = fetch_page
        self._extract_items = extract_items
        self._extract_total = extract_total
        self._page_size = page_size
        self._offset: int = 0
        self._total: int | None = None
        self._buffer: list[TItem] = []
        self._buffer_index: int = 0
        self._exhausted: bool = False

    def __aiter__(self) -> Self:
        """Return self as async iterator."""
        return self

    async def __anext__(self) -> TItem:
        """Get the next item, fetching a new page if needed."""
        # If we have items in buffer, return next one
        if self._buffer_index < len(self._buffer):
            item = self._buffer[self._buffer_index]
            self._buffer_index += 1
            return item

        # Buffer exhausted - try to fetch next page
        if self._exhausted:
            raise StopAsyncIteration

        # Fetch next page
        response = await self._fetch_page(self._page_size, self._offset)
        self._buffer = self._extract_items(response)
        self._buffer_index = 0
        self._total = self._extract_total(response)

        # Update offset for next fetch
        self._offset += len(self._buffer)

        # Check if we got any items
        if not self._buffer:
            self._exhausted = True
            raise StopAsyncIteration

        # Check if we've reached the end
        if self._offset >= self._total:
            self._exhausted = True

        # Return first item from new buffer
        item = self._buffer[self._buffer_index]
        self._buffer_index += 1
        return item

    async def pages(self) -> AsyncIterator[list[TItem]]:
        """Iterate page-by-page instead of item-by-item.

        Yields:
            List of items for each page
        """
        while not self._exhausted:
            response = await self._fetch_page(self._page_size, self._offset)
            items = self._extract_items(response)
            self._total = self._extract_total(response)

            # Update offset for next fetch
            self._offset += len(items)

            if not items:
                self._exhausted = True
                return

            if self._offset >= self._total:
                self._exhausted = True

            yield items
