import asyncio
from typing import AsyncIterator, Union
from rath.links.base import AsyncTerminatingLink
from rath.operation import GraphQLResult, Operation
from graphql import FieldNode, GraphQLSchema, OperationType, validate
import uuid
from rath.links.errors import LinkNotConnectedError


def target_from_node(node: FieldNode) -> str:
    return (
        node.alias.value if hasattr(node, "alias") and node.alias else node.name.value
    )


class AsyncMockResolver:
    def __getitem__(self, key):
        return getattr(self, f"resolve_{key}")

    def __contains__(self, key):
        return hasattr(self, f"resolve_{key}")


class AsyncStatefulMockLink(AsyncTerminatingLink):
    def __init__(
        self,
        resolver: Union[AsyncMockResolver, dict] = {},
        query_resolver: Union[AsyncMockResolver, dict] = None,
        mutation_resolver: Union[AsyncMockResolver, dict] = None,
        subscription_resolver: Union[AsyncMockResolver, dict] = None,
        schema: GraphQLSchema = None,
    ) -> None:
        self.query_resolver = query_resolver or resolver
        self.mutation_resolver = mutation_resolver or resolver
        self.subscription_resolver = subscription_resolver or resolver
        self.schema = schema
        self.connected = False
        self.futures = {}

    async def aconnect(self) -> None:
        self.connected = True
        self._inqueue = asyncio.Queue()
        self.connection_task = asyncio.create_task(self.resolving())

    async def adisconnect(self) -> None:
        self.connected = False
        self.connection_task.cancel()

        try:
            await self.connection_task
        except asyncio.CancelledError:
            pass

    async def resolving(self):
        while True:
            operation, id = await self._inqueue.get()

            resolve_futures = []

            try:
                if operation.node.operation == OperationType.QUERY:

                    for op in operation.node.selection_set.selections:
                        assert (
                            op.name.value in self.query_resolver
                        ), f"Mocked Resolver for Query '{op.name.value}' not in resolver: {self.query_resolver}  for AsyncMockLink"

                        resolve_futures.append(
                            self.query_resolver[op.name.value](operation)
                        )

                if operation.node.operation == OperationType.MUTATION:

                    for op in operation.node.selection_set.selections:
                        assert (
                            op.name.value in self.mutation_resolver
                        ), f"Mocked Resolver for Mutation {op.name.value} not in resolver: {self.mutation_resolver} for AsyncMockLink"
                        resolve_futures.append(
                            self.mutation_resolver[op.name.value](operation)
                        )

                resolved = await asyncio.gather(*resolve_futures)
                self.futures[id].set_result(
                    GraphQLResult(
                        data={
                            target_from_node(op): resolved[i]
                            for i, op in enumerate(
                                operation.node.selection_set.selections
                            )
                        }
                    )
                )

            except Exception as e:
                self.futures[id].set_exception(e)

            self._inqueue.task_done()

    async def submit(self, o, id):
        await self._inqueue.put((o, id))

    async def aquery(self, operation: Operation) -> GraphQLResult:
        if self.connected == False:
            raise LinkNotConnectedError(
                "This is a async Terminating Link, that needs to be connected before. In an async context, preferably use an async context manager. In sync call connect before."
            )

        uniqueid = uuid.uuid4()
        self.futures[uniqueid] = asyncio.Future()
        await self.submit(operation, uniqueid)
        return await self.futures[uniqueid]

    async def asubscribe(self, operation: Operation) -> AsyncIterator[GraphQLResult]:
        assert NotImplementedError("This is not implemented for now")

    def query(self, operation: Operation) -> GraphQLResult:
        raise NotImplementedError("This is a async Terminating Link")

    def subscribe(self, operation: Operation) -> GraphQLResult:
        raise NotImplementedError("This is a async Terminating Link")
