# ext/mypy/plugin.py
# Copyright (C) 2021 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php

"""
Mypy plugin for SQLAlchemy ORM.

"""
from typing import Callable
from typing import List
from typing import Optional
from typing import Tuple
from typing import Type as TypingType
from typing import Union

from mypy import nodes
from mypy.mro import calculate_mro
from mypy.mro import MroError
from mypy.nodes import Block
from mypy.nodes import ClassDef
from mypy.nodes import GDEF
from mypy.nodes import MypyFile
from mypy.nodes import NameExpr
from mypy.nodes import SymbolTable
from mypy.nodes import SymbolTableNode
from mypy.nodes import TypeInfo
from mypy.plugin import AttributeContext
from mypy.plugin import ClassDefContext
from mypy.plugin import DynamicClassDefContext
from mypy.plugin import Plugin
from mypy.plugin import SemanticAnalyzerPluginInterface
from mypy.types import get_proper_type
from mypy.types import Instance
from mypy.types import Type

from . import decl_class
from . import names
from . import util


class SQLAlchemyPlugin(Plugin):
    def get_dynamic_class_hook(
        self, fullname: str
    ) -> Optional[Callable[[DynamicClassDefContext], None]]:
        if names._type_id_for_fullname(fullname) is names.DECLARATIVE_BASE:
            return _dynamic_class_hook
        return None

    def get_base_class_hook(
        self, fullname: str
    ) -> Optional[Callable[[ClassDefContext], None]]:

        # kind of a strange relationship between get_metaclass_hook()
        # and get_base_class_hook().  the former doesn't fire off for
        # subclasses.   but then you can just check it here from the "base"
        # and get the same effect.
        sym = self.lookup_fully_qualified(fullname)

        if (
            sym
            and isinstance(sym.node, TypeInfo)
            and sym.node.metaclass_type
            and names._type_id_for_named_node(sym.node.metaclass_type.type)
            is names.DECLARATIVE_META
        ):
            return _base_cls_hook
        return None

    def get_class_decorator_hook(
        self, fullname: str
    ) -> Optional[Callable[[ClassDefContext], None]]:

        sym = self.lookup_fully_qualified(fullname)

        if sym is not None and sym.node is not None:
            type_id = names._type_id_for_named_node(sym.node)
            if type_id is names.MAPPED_DECORATOR:
                return _cls_decorator_hook
            elif type_id in (
                names.AS_DECLARATIVE,
                names.AS_DECLARATIVE_BASE,
            ):
                return _base_cls_decorator_hook
            elif type_id is names.DECLARATIVE_MIXIN:
                return _declarative_mixin_hook

        return None

    def get_customize_class_mro_hook(
        self, fullname: str
    ) -> Optional[Callable[[ClassDefContext], None]]:
        return _fill_in_decorators

    def get_attribute_hook(
        self, fullname: str
    ) -> Optional[Callable[[AttributeContext], Type]]:
        if fullname.startswith(
            "sqlalchemy.orm.attributes.QueryableAttribute."
        ):
            return _queryable_getattr_hook
        return None

    def get_additional_deps(
        self, file: MypyFile
    ) -> List[Tuple[int, str, int]]:
        return [
            (10, "sqlalchemy.orm.attributes", -1),
            (10, "sqlalchemy.orm.decl_api", -1),
        ]


def plugin(version: str) -> TypingType[SQLAlchemyPlugin]:
    return SQLAlchemyPlugin


def _queryable_getattr_hook(ctx: AttributeContext) -> Type:
    # how do I....tell it it has no attribute of a certain name?
    # can't find any Type that seems to match that
    return ctx.default_attr_type


def _fill_in_decorators(ctx: ClassDefContext) -> None:
    for decorator in ctx.cls.decorators:
        # set the ".fullname" attribute of a class decorator
        # that is a MemberExpr.   This causes the logic in
        # semanal.py->apply_class_plugin_hooks to invoke the
        # get_class_decorator_hook for our "registry.map_class()"
        # and "registry.as_declarative_base()" methods.
        # this seems like a bug in mypy that these decorators are otherwise
        # skipped.

        if (
            isinstance(decorator, nodes.CallExpr)
            and isinstance(decorator.callee, nodes.MemberExpr)
            and decorator.callee.name == "as_declarative_base"
        ):
            target = decorator.callee
        elif (
            isinstance(decorator, nodes.MemberExpr)
            and decorator.name == "mapped"
        ):
            target = decorator
        else:
            continue

        assert isinstance(target.expr, NameExpr)
        sym = ctx.api.lookup_qualified(
            target.expr.name, target, suppress_errors=True
        )
        if sym and sym.node:
            sym_type = get_proper_type(sym.type)
            if isinstance(sym_type, Instance):
                target.fullname = f"{sym_type.type.fullname}.{target.name}"
            else:
                # if the registry is in the same file as where the
                # decorator is used, it might not have semantic
                # symbols applied and we can't get a fully qualified
                # name or an inferred type, so we are actually going to
                # flag an error in this case that they need to annotate
                # it.  The "registry" is declared just
                # once (or few times), so they have to just not use
                # type inference for its assignment in this one case.
                util.fail(
                    ctx.api,
                    "Class decorator called %s(), but we can't "
                    "tell if it's from an ORM registry.  Please "
                    "annotate the registry assignment, e.g. "
                    "my_registry: registry = registry()" % target.name,
                    sym.node,
                )


def _add_globals(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> None:
    """Add __sa_DeclarativeMeta and __sa_Mapped symbol to the global space
    for all class defs

    """

    util.add_global(
        ctx,
        "sqlalchemy.orm.decl_api",
        "DeclarativeMeta",
        "__sa_DeclarativeMeta",
    )

    util.add_global(ctx, "sqlalchemy.orm.attributes", "Mapped", "__sa_Mapped")


def _cls_metadata_hook(ctx: ClassDefContext) -> None:
    _add_globals(ctx)
    decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)


def _base_cls_hook(ctx: ClassDefContext) -> None:
    _add_globals(ctx)
    decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)


def _declarative_mixin_hook(ctx: ClassDefContext) -> None:
    _add_globals(ctx)
    decl_class._scan_declarative_assignments_and_apply_types(
        ctx.cls, ctx.api, is_mixin_scan=True
    )


def _cls_decorator_hook(ctx: ClassDefContext) -> None:
    _add_globals(ctx)
    assert isinstance(ctx.reason, nodes.MemberExpr)
    expr = ctx.reason.expr

    assert isinstance(expr, nodes.RefExpr) and isinstance(expr.node, nodes.Var)

    node_type = get_proper_type(expr.node.type)

    assert (
        isinstance(node_type, Instance)
        and names._type_id_for_named_node(node_type.type) is names.REGISTRY
    )

    decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)


def _base_cls_decorator_hook(ctx: ClassDefContext) -> None:
    _add_globals(ctx)

    cls = ctx.cls

    _make_declarative_meta(ctx.api, cls)

    decl_class._scan_declarative_assignments_and_apply_types(
        cls, ctx.api, is_mixin_scan=True
    )


def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
    """Generate a declarative Base class when the declarative_base() function
    is encountered."""

    _add_globals(ctx)

    cls = ClassDef(ctx.name, Block([]))
    cls.fullname = ctx.api.qualified_name(ctx.name)

    info = TypeInfo(SymbolTable(), cls, ctx.api.cur_mod_id)
    cls.info = info
    _make_declarative_meta(ctx.api, cls)

    cls_arg = util._get_callexpr_kwarg(ctx.call, "cls", expr_types=(NameExpr,))
    if cls_arg is not None and isinstance(cls_arg.node, TypeInfo):
        decl_class._scan_declarative_assignments_and_apply_types(
            cls_arg.node.defn, ctx.api, is_mixin_scan=True
        )
        info.bases = [Instance(cls_arg.node, [])]
    else:
        obj = ctx.api.named_type("__builtins__.object")

        info.bases = [obj]

    try:
        calculate_mro(info)
    except MroError:
        util.fail(
            ctx.api, "Not able to calculate MRO for declarative base", ctx.call
        )
        obj = ctx.api.named_type("__builtins__.object")
        info.bases = [obj]
        info.fallback_to_any = True

    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))


def _make_declarative_meta(
    api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
) -> None:

    declarative_meta_name: NameExpr = NameExpr("__sa_DeclarativeMeta")
    declarative_meta_name.kind = GDEF
    declarative_meta_name.fullname = "sqlalchemy.orm.decl_api.DeclarativeMeta"

    # installed by _add_globals
    sym = api.lookup_qualified("__sa_DeclarativeMeta", target_cls)

    assert sym is not None and isinstance(sym.node, nodes.TypeInfo)

    declarative_meta_typeinfo = sym.node
    declarative_meta_name.node = declarative_meta_typeinfo

    target_cls.metaclass = declarative_meta_name

    declarative_meta_instance = Instance(declarative_meta_typeinfo, [])

    info = target_cls.info
    info.declared_metaclass = info.metaclass_type = declarative_meta_instance