from aiodataloader import DataLoader
from ariadne import ObjectType, QueryType, make_executable_schema
from ariadne.asgi import GraphQL
from fastapi import FastAPI, Request
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from sqlalchemy_relay_pagination import Connection, paginate, paginate_many
from sqlalchemy_relay_pagination.graphql import extract_requested_fields

# Re-use models from top-level conftest
from tests.conftest import Post, User

SDL = """
    type Query {
        users(first: Int, last: Int, after: String, before: String): UserConnection!
    }

    type UserConnection {
        edges: [UserEdge!]!
        pageInfo: PageInfo!
    }

    type UserEdge {
        node: User!
        cursor: String!
    }

    type PageInfo {
        hasNextPage: Boolean!
        hasPreviousPage: Boolean!
        startCursor: String
        endCursor: String
        totalCount: Int
    }

    type User {
        id: Int!
        name: String!
        posts(first: Int, last: Int): PostConnection!
    }

    type PostConnection {
        edges: [PostEdge!]!
        pageInfo: PageInfo!
    }

    type PostEdge {
        node: Post!
        cursor: String!
    }

    type Post {
        id: Int!
        title: String!
    }
"""


query_type = QueryType()
user_type = ObjectType("User")


@query_type.field("users")
async def resolve_users(_, info, first=None, last=None, after=None, before=None):
    session: AsyncSession = info.context["session"]

    requested = extract_requested_fields(info)
    query = select(User)
    conn = await paginate(
        session,
        query,
        order_by=User.id,
        first=first,
        last=last,
        after=after,
        before=before,
        requested_fields=requested,
    )

    return conn


class UserPostsLoader(DataLoader):
    """Batches User.posts resolvers within a single GraphQL request."""

    def __init__(self, session: AsyncSession, first, last, fields):
        super().__init__()
        self._session = session
        self._first = first
        self._last = last
        self._fields = fields

    async def batch_load_fn(self, user_ids):
        results = await paginate_many(
            self._session,
            select(Post),
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=list(user_ids),
            first=self._first,
            last=self._last,
            requested_fields=self._fields,
        )
        return [results[uid] for uid in user_ids]


@user_type.field("posts")
async def resolve_user_posts(user: User, info, first: int | None = None, last: int | None = None):
    session: AsyncSession = info.context["session"]
    loaders: dict[tuple[int | None, int | None], DataLoader[int, Connection[Post]]] = info.context["loaders"]

    loader_key = (first, last)
    if loader_key not in loaders:
        fields = extract_requested_fields(info)
        loaders[loader_key] = UserPostsLoader(session, first, last, fields)

    loader = loaders[loader_key]
    conn = await loader.load(user.id)
    return conn


schema = make_executable_schema(SDL, query_type, user_type)


def create_app(session_factory) -> FastAPI:
    app = FastAPI()

    @app.middleware("http")
    async def session_middleware(request: Request, call_next):
        async with session_factory() as session:
            request.state.session = session
            return await call_next(request)

    async def get_context(request, data):
        return {"session": request.state.session, "request": request, "loaders": {}}

    graphql_app = GraphQL(schema, context_value=get_context)
    app.mount("/graphql", graphql_app)
    return app
