from mypy.plugin import Plugin
from mypy.nodes import (
    ClassDef, Block, TypeInfo, SymbolTable, SymbolTableNode, GDEF, Var
)
from mypy.types import Instance

DECL_BASES = set()

class DynPlugin(Plugin):
    def get_dynamic_class_hook(self, fullname):
        if fullname == 'mod.declarative_base':
            return add_info_hook
        return None

    def get_base_class_hook(self, fullname: str):
        if fullname in DECL_BASES:
            return replace_col_hook
        return None

def add_info_hook(ctx):
    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):
    info = ctx.cls.info
    for sym in info.names.values():
        node = sym.node
        if isinstance(node, Var) and isinstance(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):
    return DynPlugin
