from __future__ import annotations

import contextlib
from contextlib import suppress
from typing import TYPE_CHECKING

from django.core.exceptions import FieldDoesNotExist
from django.db.models import Field, ForeignKey, Model
from graphene import Connection, ObjectType, PageInfo
from graphene.relay.node import AbstractNode
from graphene.types.definitions import GrapheneInterfaceType, GrapheneObjectType, GrapheneUnionType
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,
    FragmentDefinitionNode,
    FragmentSpreadNode,
    GraphQLField,
    GraphQLOutputType,
    GraphQLSchema,
    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, GQLInfo, ModelField, Optional, ToManyField, ToOneField, TypeGuard, Union, overload

if TYPE_CHECKING:
    from graphene.types.base import BaseOptions

__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: Optional[type[Model]] = 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)

        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 all(field in field_type.fields for field in ["node", "cursor"]) and field_type.name.endswith("Edge")


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:
    return issubclass(getattr(getattr(graphql_field.resolve, "func", None), "__self__", type(None)), AbstractNode)


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: Optional[GrapheneObjectType] = 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: Optional[GrapheneObjectType] = 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) -> Optional[ModelField]:
    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
