from __future__ import annotations

import contextlib
from contextlib import suppress
from types import MethodType
from typing import TYPE_CHECKING

from django.core.exceptions import FieldDoesNotExist
from django.db.models import ForeignKey
from graphene import Connection, ObjectType, PageInfo
from graphene.relay.node import AbstractNode
from graphene.types.definitions import GrapheneObjectType, GrapheneUnionType
from graphene.types.union import Union as GrapheneUnion
from graphene.utils.str_converters import to_snake_case
from graphene_django import DjangoObjectType
from graphene_django.registry import get_global_registry
from graphql import FieldNode, FragmentSpreadNode, InlineFragmentNode, SelectionNode, Undefined
from graphql.execution.execute import get_field_def

from .errors import OptimizerError
from .settings import optimizer_settings
from .typing import GRAPHQL_BUILTIN, Union, overload

if TYPE_CHECKING:
    from django.db.models import Field, Model
    from graphene.types.base import BaseOptions
    from graphene.types.definitions import GrapheneInterfaceType
    from graphql import FragmentDefinitionNode, GraphQLField, GraphQLOutputType, GraphQLSchema

    from .typing import GQLInfo, ModelField, ToManyField, ToOneField, TypeGuard

__all__ = [
    "GraphQLASTWalker",
]


GrapheneType = Union[GrapheneObjectType, GrapheneUnionType]
Selections = tuple[SelectionNode, ...]


class GraphQLASTWalker:
    """Class for walking the GraphQL AST and handling the different nodes."""

    def __init__(self, info: GQLInfo, model: type[Model] | None = None) -> None:
        self.info = info
        self.complexity: int = 0
        self.model: type[Model] = model

    def increase_complexity(self) -> None:
        self.complexity += 1

    def run(self) -> None:
        selections: Selections = self.info.field_nodes  # type: ignore[assignment]
        field_type: GrapheneObjectType = self.info.parent_type  # type: ignore[assignment]
        return self.handle_selections(field_type, selections)

    def handle_selections(self, field_type: GrapheneType, selections: Selections) -> None:
        for selection in selections:
            if isinstance(selection, FieldNode):
                self.handle_field_node(field_type, selection)

            elif isinstance(selection, FragmentSpreadNode):
                self.handle_fragment_spread(field_type, selection)

            elif isinstance(selection, InlineFragmentNode):
                self.handle_inline_fragment(field_type, selection)

            else:  # pragma: no cover
                msg = f"Unhandled selection node: '{selection}'"
                raise OptimizerError(msg)

    def handle_field_node(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        graphene_type: type[ObjectType] = field_type.graphene_type

        if self.info.parent_type == field_type:
            return self.handle_query_class(field_type, field_node)

        if issubclass(graphene_type, Connection):
            return self.handle_connection(field_type, field_node)

        if is_edge(field_type):
            return self.handle_edge(field_type, field_node)

        if issubclass(graphene_type, PageInfo):  # pragma: no cover
            return self.handle_page_info(field_type, field_node)

        if issubclass(graphene_type, ObjectType):
            return self.handle_object_type(field_type, field_node)

        if issubclass(graphene_type, AbstractNode):
            return self.handle_abstract_node(field_type, field_node)

        if issubclass(graphene_type, GrapheneUnion):
            return self.handle_abstract_node(field_type, field_node)

        msg = f"Unhandled graphene type: '{graphene_type}'"  # pragma: no cover
        raise OptimizerError(msg)  # pragma: no cover

    def handle_query_class(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        graphene_type = self.get_graphene_type(field_type, field_node)
        selections = get_selections(field_node)
        return self.handle_selections(graphene_type, selections)

    def handle_object_type(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        field_name = to_snake_case(field_node.name.value)
        if is_graphql_builtin(field_name):
            return self.handle_graphql_builtin(field_type, field_node)

        graphene_type: type[ObjectType] = field_type.graphene_type

        if issubclass(graphene_type, DjangoObjectType):
            return self.handle_model_field(field_type, field_node, field_name)
        return self.handle_plain_object_type(field_type, field_node)

    def handle_abstract_node(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        graphene_type = get_global_registry().get_type_for_model(self.model._meta.concrete_model)
        field_type = self.info.schema.get_type(graphene_type._meta.name)
        return self.handle_object_type(field_type, field_node)

    def handle_graphql_builtin(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None: ...

    def handle_plain_object_type(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None: ...

    def handle_model_field(self, field_type: GrapheneObjectType, field_node: FieldNode, field_name: str) -> None:
        model: type[Model] = field_type.graphene_type._meta.model
        field = get_model_field(model, field_name)

        if field is None:
            with self.use_model(model):
                return self.handle_custom_field(field_type, field_node)

        if not field.is_relation or is_foreign_key_id(field, field_node):
            with self.use_model(model):
                return self.handle_normal_field(field_type, field_node, field)

        if is_to_one(field):
            related_model = get_related_model(field, model)
            with self.use_model(field.model):
                return self.handle_to_one_field(field_type, field_node, field, related_model)

        if is_to_many(field):
            related_model = get_related_model(field, model)
            with self.use_model(model):
                return self.handle_to_many_field(field_type, field_node, field, related_model)

        msg = f"Unhandled field: '{field.name}'"  # pragma: no cover
        raise OptimizerError(msg)  # pragma: no cover

    def handle_custom_field(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None: ...

    def handle_normal_field(self, field_type: GrapheneObjectType, field_node: FieldNode, field: Field) -> None: ...

    def handle_to_one_field(
        self,
        field_type: GrapheneObjectType,
        field_node: FieldNode,
        related_field: ToOneField,
        related_model: type[Model] | None,
    ) -> None:
        graphene_type = self.get_graphene_type(field_type, field_node)
        selections = get_selections(field_node)
        self.increase_complexity()
        return self.handle_selections(graphene_type, selections)

    def handle_to_many_field(
        self,
        field_type: GrapheneObjectType,
        field_node: FieldNode,
        related_field: ToManyField,
        related_model: type[Model] | None,
    ) -> None:
        graphene_type = self.get_graphene_type(field_type, field_node)
        selections = get_selections(field_node)
        self.increase_complexity()
        return self.handle_selections(graphene_type, selections)

    def handle_connection(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        if field_node.name.value == optimizer_settings.TOTAL_COUNT_FIELD:
            return self.handle_total_count(field_type, field_node)

        graphene_type = self.get_graphene_type(field_type, field_node)
        selections = get_selections(field_node)
        return self.handle_selections(graphene_type, selections)

    def handle_edge(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None:
        graphene_type = self.get_graphene_type(field_type, field_node)
        selections = get_selections(field_node)
        return self.handle_selections(graphene_type, selections)

    def handle_total_count(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None: ...

    def handle_page_info(self, field_type: GrapheneObjectType, field_node: FieldNode) -> None: ...

    def handle_fragment_spread(self, field_type: GrapheneObjectType, fragment_spread: FragmentSpreadNode) -> None:
        name = fragment_spread.name.value
        fragment_definition = self.info.fragments[name]
        selections = get_selections(fragment_definition)
        return self.handle_selections(field_type, selections)

    def handle_inline_fragment(
        self,
        field_type: GrapheneUnionType | GrapheneInterfaceType,
        inline_fragment: InlineFragmentNode,
    ) -> None:
        fragment_type = get_fragment_type(field_type, inline_fragment, self.info.schema)
        graphene_options: BaseOptions = fragment_type.graphene_type._meta
        fragment_model: type[Model] | Undefined = getattr(graphene_options, "model", Undefined)

        if fragment_model != self.model:
            return None

        selections = get_selections(inline_fragment)
        return self.handle_selections(fragment_type, selections)

    def get_graphene_type(self, field_type: GrapheneObjectType, field_node: FieldNode) -> GrapheneType:
        graphql_field = get_field_def(self.info.schema, field_type, field_node)
        return get_underlying_type(graphql_field.type)

    def get_field_name(self, field_node: FieldNode) -> str:
        alias = getattr(field_node.alias, "value", None)
        return alias or to_snake_case(field_node.name.value)

    def get_related_field_name(self, related_field: ToOneField | ToManyField) -> str:  # pragma: no cover
        if hasattr(related_field, "cache_name"):  # New in Django 5.1
            return related_field.cache_name or related_field.name
        return related_field.get_cache_name() or related_field.name

    @contextlib.contextmanager
    def use_model(self, model: type[Model]) -> GraphQLASTWalker:
        orig_model = self.model
        try:
            self.model = model
            yield
        finally:
            self.model = orig_model


@overload
def get_underlying_type(
    field_type: type[GraphQLOutputType],
) -> type[Union[DjangoObjectType, GrapheneObjectType]]: ...  # pragma: no cover


@overload
def get_underlying_type(
    field_type: GraphQLOutputType,
) -> Union[DjangoObjectType, GrapheneObjectType]: ...  # pragma: no cover


def get_underlying_type(field_type):
    while hasattr(field_type, "of_type"):
        field_type = field_type.of_type
    return field_type


def get_selections(field_node: Union[FieldNode, FragmentDefinitionNode, InlineFragmentNode]) -> Selections:
    return getattr(field_node.selection_set, "selections", ())


def is_edge(field_type: GrapheneObjectType) -> bool:
    # Edge-classes are created by `graphene.relay.connection.get_edge_class`,
    # which means that we cannot check against the EdgeBase class directly.
    return (
        isinstance(field_type, GrapheneObjectType)
        and field_type.name.endswith("Edge")
        and "cursor" in field_type.fields
        and "node" in field_type.fields
    )


def is_connection(graphql_field: Union[GraphQLField, GrapheneObjectType]) -> bool:
    return issubclass(getattr(graphql_field, "graphene_type", type(None)), Connection)


def is_node(graphql_field: GraphQLField) -> bool:
    # Assumes node is resolved using 'graphene.relay.node.Node.node_resolver'
    resolver = getattr(graphql_field.resolve, "func", None)
    if isinstance(resolver, MethodType):
        cls: type = resolver.__self__  # type: ignore[assignment]
        return isinstance(cls, type) and issubclass(cls, AbstractNode)
    return False


def is_graphql_builtin(field_name: str) -> bool:
    return field_name.lower() in GRAPHQL_BUILTIN


def is_foreign_key_id(field: Field, field_node: FieldNode) -> bool:
    return isinstance(field, ForeignKey) and field.get_attname() == to_snake_case(field_node.name.value)


def is_to_many(field: Field) -> TypeGuard[ToManyField]:
    return bool(field.one_to_many or field.many_to_many)


def is_to_one(field: Field) -> TypeGuard[ToOneField]:
    return bool(field.many_to_one or field.one_to_one)


def get_fragment_type(
    field_type: GrapheneUnionType | GrapheneInterfaceType,
    inline_fragment: InlineFragmentNode,
    schema: GraphQLSchema,
) -> GrapheneObjectType:
    fragment_type_name = inline_fragment.type_condition.name.value

    # For unions, fetch the type from in the union.
    if isinstance(field_type, GrapheneUnionType):
        gen = (t for t in field_type.types if t.name == fragment_type_name)
        fragment_type: GrapheneObjectType | None = next(gen, None)
        if fragment_type is None:  # pragma: no cover
            msg = f"Fragment type '{fragment_type_name}' not found in union '{field_type}'"
            raise OptimizerError(msg)

    # For interfaces, fetch the type from in the schema.
    else:
        fragment_type: GrapheneObjectType | None = schema.get_type(fragment_type_name)
        if fragment_type is None:  # pragma: no cover
            msg = f"Fragment type '{fragment_type_name}' not found in schema."
            raise OptimizerError(msg)

    return fragment_type


def get_related_model(related_field: Union[ToOneField, ToManyField], model: type[Model]) -> type[Model] | None:
    """
    Get the related model for a field.
    Note: For generic foreign keys, the related model is unknown (=None).
    """
    related_model = related_field.related_model
    if related_model == "self":  # pragma: no cover
        return model
    return related_model  # type: ignore[return-value]


def get_model_field(model: type[Model], field_name: str) -> ModelField | None:
    if field_name == "pk":
        model_field: ModelField = model._meta.pk
        return model_field

    with suppress(FieldDoesNotExist):
        model_field: ModelField = model._meta.get_field(field_name)
        return model_field

    # Field might be a reverse many-related field without `related_name`, in which case
    # the `model._meta.fields_map` will store the relation without the "_set" suffix.
    if field_name.endswith("_set"):
        with suppress(FieldDoesNotExist):
            model_field: ModelField = model._meta.get_field(field_name.removesuffix("_set"))
            if is_to_many(model_field):
                return model_field

    return None
