from __future__ import annotations

from typing import Callable

from mypy.nodes import NameExpr
from mypy.plugin import FunctionContext, Plugin
from mypy.types import Instance, NoneType, Type, UnionType, get_proper_type


class AttrPlugin(Plugin):
    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
        if fullname.startswith("mod.Attr"):
            return attr_hook
        return None


def attr_hook(ctx: FunctionContext) -> Type:
    default = get_proper_type(ctx.default_return_type)
    assert isinstance(default, Instance)
    if default.type.fullname == "mod.Attr":
        attr_base = default
    else:
        attr_base = None
    for base in default.type.bases:
        if base.type.fullname == "mod.Attr":
            attr_base = base
            break
    assert attr_base is not None
    last_arg_exprs = ctx.args[-1]
    if any(isinstance(expr, NameExpr) and expr.name == "True" for expr in last_arg_exprs):
        return attr_base
    assert len(attr_base.args) == 1
    arg_type = attr_base.args[0]
    return Instance(
        attr_base.type,
        [UnionType([arg_type, NoneType()])],
        line=default.line,
        column=default.column,
    )


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