import typing
from label_studio_sdk.core.pagination import SyncPager, AsyncPager, T
from label_studio_sdk.core.api_error import ApiError

# This is a custom extension of the autogenerated SyncPager and AsyncPager classes
# that works with the Label Studio SDK's default pagination behavior
# that throws 404 errors at the end of the pagination.


class SyncPagerExt(SyncPager[T, typing.Any], typing.Generic[T]):
    @classmethod
    def from_sync_pager(cls, sync_pager: SyncPager) -> "SyncPagerExt[T]":
        # Minimal compatibility with newer Fern pagers:
        # SyncPager is a frozen dataclass and now *requires* the `response` field.
        return cls(
            get_next=sync_pager.get_next,
            has_next=sync_pager.has_next,
            items=sync_pager.items,
            response=sync_pager.response,
        )
 
    def __iter__(self) -> typing.Iterator[T]:  # type: ignore
        # Extends the iterator to catch 404 errors at the end of the pagination
        try:
            for item in super().__iter__():
                yield item
        except ApiError as exc:
            if exc.status_code == 404:
                return
            raise


class AsyncPagerExt(AsyncPager[T, typing.Any], typing.Generic[T]):
    @classmethod
    async def from_async_pager(cls, async_pager: AsyncPager) -> "AsyncPagerExt[T]":
        return cls(
            get_next=async_pager.get_next,
            has_next=async_pager.has_next,
            items=async_pager.items,
            response=async_pager.response,
        )
 
    async def __aiter__(self) -> typing.AsyncIterator[T]:  # type: ignore
        # Extends the iterator to catch 404 errors at the end of the pagination
        try:
            async for item in super().__aiter__():
                yield item
        except ApiError as exc:
            if exc.status_code == 404:
                return
            raise
 
    async def __anext__(self) -> T:
        try:
            return await super().__anext__()
        except ApiError as exc:
            if exc.status_code == 404:
                raise StopAsyncIteration
            raise