from __future__ import annotations

from typing import Callable

from mypy.nodes import GDEF, Block, ClassDef, SymbolTable, SymbolTableNode, TypeInfo, Var
from mypy.plugin import ClassDefContext, DynamicClassDefContext, Plugin
from mypy.types import Instance, get_proper_type

DECL_BASES: set[str] = set()


class DynPlugin(Plugin):
    def get_dynamic_class_hook(
        self, fullname: str
    ) -> Callable[[DynamicClassDefContext], None] | None:
        if fullname == "mod.declarative_base":
            return add_info_hook
        return None

    def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None:
        if fullname in DECL_BASES:
            return replace_col_hook
        return None


def add_info_hook(ctx: DynamicClassDefContext) -> None:
    class_def = ClassDef(ctx.name, Block([]))
    class_def.fullname = ctx.api.qualified_name(ctx.name)

    info = TypeInfo(SymbolTable(), class_def, ctx.api.cur_mod_id)
    class_def.info = info
    obj = ctx.api.named_type("builtins.object")
    info.mro = [info, obj.type]
    info.bases = [obj]
    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
    DECL_BASES.add(class_def.fullname)


def replace_col_hook(ctx: ClassDefContext) -> None:
    info = ctx.cls.info
    for sym in info.names.values():
        node = sym.node
        if isinstance(node, Var) and isinstance(
            (node_type := get_proper_type(node.type)), Instance
        ):
            if node_type.type.fullname == "mod.Column":
                new_sym = ctx.api.lookup_fully_qualified_or_none("mod.Instr")
                if new_sym:
                    new_info = new_sym.node
                    assert isinstance(new_info, TypeInfo)
                    node.type = Instance(
                        new_info, node_type.args, node_type.line, node_type.column
                    )


def plugin(version: str) -> type[DynPlugin]:
    return DynPlugin