import math
from typing import Any, Generic, List, TypeVar

from django.db.models import QuerySet
from ninja import Schema
from ninja.pagination import PageNumberPagination

T = TypeVar("T")


class EnhancedPageNumberPagination(PageNumberPagination):
    class Output(Schema, Generic[T]):
        items: List[T]
        page: int
        page_size: int
        total_items: int
        total_pages: int

    def paginate_queryset(
        self,
        queryset: QuerySet,
        pagination: PageNumberPagination.Input,
        **params: Any,
    ) -> Any:
        page_size = self._get_page_size(pagination.page_size)
        offset = (pagination.page - 1) * page_size
        total_items = self._items_count(queryset)
        total_pages = math.ceil(total_items / page_size) if page_size > 0 else 0

        return {
            "items": queryset[offset : offset + page_size],
            "page": pagination.page,
            "page_size": page_size,
            "total_items": total_items,
            "total_pages": total_pages,
        }

    async def apaginate_queryset(
        self,
        queryset: QuerySet,
        pagination: PageNumberPagination.Input,
        **params: Any,
    ) -> Any:
        page_size = self._get_page_size(pagination.page_size)
        offset = (pagination.page - 1) * page_size
        total_items = await self._aitems_count(queryset)
        total_pages = math.ceil(total_items / page_size) if page_size > 0 else 0

        if isinstance(queryset, QuerySet):
            items = [obj async for obj in queryset[offset : offset + page_size]]
        else:
            items = queryset[offset : offset + page_size]

        return {
            "items": items,
            "page": pagination.page,
            "page_size": page_size,
            "total_items": total_items,
            "total_pages": total_pages,
        }