import strawberry
from strawberry.relay.utils import from_base64, to_base64

import strawberry_django
from strawberry_django import relay

@strawberry.type
class CustomConnection(relay.ListConnectionWithTotalCount[strawberry.relay.NodeType]):
    """
    A strawberry connection to count the number of query results
    """

    # Adding offset argument to custom connection
    @classmethod
    def resolve_connection(
        cls,
        nodes: NodeIterableType[strawberry.relay.NodeType],
        *,
        info: Info,
        before: Optional[str] = None,
        after: Optional[str] = None,
        first: Optional[int] = None,
        last: Optional[int] = None,
        offset: Optional[int] = None,
        **kwargs: Any,
    ) -> AwaitableOrValue[Self]:

        # This implemntation is based on the graphene
        # implementation of first/offset pagination
        if offset:
            if after:
                offset += from_base64(after) + 1
            # input offset starts at 1 while the offset starts at 0
            after = to_base64("arrayconnection", offset - 1)

        conn = super().resolve_connection(
            nodes,
            info=info,
            before=before,
            after=after,
            first=first,
            last=last,
            **kwargs,
        )

        if inspect.isawaitable(conn):

            async def wrapper():
                resolved = await conn
                resolved.nodes = nodes
                return resolved

            return wrapper()

        conn = cast(Self, conn)
        conn.nodes = nodes
        return conn

    @strawberry.field
    def foos(
        self,
        info: strawberry.types.Info,
        before: Optional[str] = None,
        after: Optional[str] = None,
        first: Optional[int] = None,
        last: Optional[int] = None,
        offset: Optional[int] = None,
    ) -> Optional[CustomConnection[FooType]]:
        # Resolver logic here returning queryset qs
        return CustomConnection[FooType].resolve_connection(info=info, nodes=qs, offset=offset, first=first, last=last, after=after, before=before)