"""Maintain a mapping from mypy concepts to IR/compiled concepts."""

from __future__ import annotations

from mypy.nodes import ARG_STAR, ARG_STAR2, GDEF, ArgKind, FuncDef, RefExpr, SymbolNode, TypeInfo
from mypy.types import (
    AnyType,
    CallableType,
    Instance,
    LiteralType,
    NoneTyp,
    Overloaded,
    PartialType,
    TupleType,
    Type,
    TypedDictType,
    TypeType,
    TypeVarType,
    UnboundType,
    UninhabitedType,
    UnionType,
    get_proper_type,
)
from mypyc.ir.class_ir import ClassIR
from mypyc.ir.func_ir import FuncDecl, FuncSignature, RuntimeArg
from mypyc.ir.rtypes import (
    RInstance,
    RTuple,
    RType,
    RUnion,
    bool_rprimitive,
    bytes_rprimitive,
    dict_rprimitive,
    float_rprimitive,
    int32_rprimitive,
    int64_rprimitive,
    int_rprimitive,
    list_rprimitive,
    none_rprimitive,
    object_rprimitive,
    range_rprimitive,
    set_rprimitive,
    str_rprimitive,
    tuple_rprimitive,
)


class Mapper:
    """Keep track of mappings from mypy concepts to IR concepts.

    For example, we keep track of how the mypy TypeInfos of compiled
    classes map to class IR objects.

    This state is shared across all modules being compiled in all
    compilation groups.
    """

    def __init__(self, group_map: dict[str, str | None]) -> None:
        self.group_map = group_map
        self.type_to_ir: dict[TypeInfo, ClassIR] = {}
        self.func_to_decl: dict[SymbolNode, FuncDecl] = {}

    def type_to_rtype(self, typ: Type | None) -> RType:
        if typ is None:
            return object_rprimitive

        typ = get_proper_type(typ)
        if isinstance(typ, Instance):
            if typ.type.fullname == "builtins.int":
                return int_rprimitive
            elif typ.type.fullname == "builtins.float":
                return float_rprimitive
            elif typ.type.fullname == "builtins.bool":
                return bool_rprimitive
            elif typ.type.fullname == "builtins.str":
                return str_rprimitive
            elif typ.type.fullname == "builtins.bytes":
                return bytes_rprimitive
            elif typ.type.fullname == "builtins.list":
                return list_rprimitive
            # Dict subclasses are at least somewhat common and we
            # specifically support them, so make sure that dict operations
            # get optimized on them.
            elif any(cls.fullname == "builtins.dict" for cls in typ.type.mro):
                return dict_rprimitive
            elif typ.type.fullname == "builtins.set":
                return set_rprimitive
            elif typ.type.fullname == "builtins.tuple":
                return tuple_rprimitive  # Varying-length tuple
            elif typ.type.fullname == "builtins.range":
                return range_rprimitive
            elif typ.type in self.type_to_ir:
                inst = RInstance(self.type_to_ir[typ.type])
                # Treat protocols as Union[protocol, object], so that we can do fast
                # method calls in the cases where the protocol is explicitly inherited from
                # and fall back to generic operations when it isn't.
                if typ.type.is_protocol:
                    return RUnion([inst, object_rprimitive])
                else:
                    return inst
            elif typ.type.fullname == "mypy_extensions.i64":
                return int64_rprimitive
            elif typ.type.fullname == "mypy_extensions.i32":
                return int32_rprimitive
            else:
                return object_rprimitive
        elif isinstance(typ, TupleType):
            # Use our unboxed tuples for raw tuples but fall back to
            # being boxed for NamedTuple.
            if typ.partial_fallback.type.fullname == "builtins.tuple":
                return RTuple([self.type_to_rtype(t) for t in typ.items])
            else:
                return tuple_rprimitive
        elif isinstance(typ, CallableType):
            return object_rprimitive
        elif isinstance(typ, NoneTyp):
            return none_rprimitive
        elif isinstance(typ, UnionType):
            return RUnion([self.type_to_rtype(item) for item in typ.items])
        elif isinstance(typ, AnyType):
            return object_rprimitive
        elif isinstance(typ, TypeType):
            return object_rprimitive
        elif isinstance(typ, TypeVarType):
            # Erase type variable to upper bound.
            # TODO: Erase to union if object has value restriction?
            return self.type_to_rtype(typ.upper_bound)
        elif isinstance(typ, PartialType):
            assert typ.var.type is not None
            return self.type_to_rtype(typ.var.type)
        elif isinstance(typ, Overloaded):
            return object_rprimitive
        elif isinstance(typ, TypedDictType):
            return dict_rprimitive
        elif isinstance(typ, LiteralType):
            return self.type_to_rtype(typ.fallback)
        elif isinstance(typ, (UninhabitedType, UnboundType)):
            # Sure, whatever!
            return object_rprimitive

        # I think we've covered everything that is supposed to
        # actually show up, so anything else is a bug somewhere.
        assert False, "unexpected type %s" % type(typ)

    def get_arg_rtype(self, typ: Type, kind: ArgKind) -> RType:
        if kind == ARG_STAR:
            return tuple_rprimitive
        elif kind == ARG_STAR2:
            return dict_rprimitive
        else:
            return self.type_to_rtype(typ)

    def fdef_to_sig(self, fdef: FuncDef) -> FuncSignature:
        if isinstance(fdef.type, CallableType):
            arg_types = [
                self.get_arg_rtype(typ, kind)
                for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds)
            ]
            arg_pos_onlys = [name is None for name in fdef.type.arg_names]
            ret = self.type_to_rtype(fdef.type.ret_type)
        else:
            # Handle unannotated functions
            arg_types = [object_rprimitive for _ in fdef.arguments]
            arg_pos_onlys = [arg.pos_only for arg in fdef.arguments]
            # We at least know the return type for __init__ methods will be None.
            is_init_method = fdef.name == "__init__" and bool(fdef.info)
            if is_init_method:
                ret = none_rprimitive
            else:
                ret = object_rprimitive

        # mypyc FuncSignatures (unlike mypy types) want to have a name
        # present even when the argument is position only, since it is
        # the sole way that FuncDecl arguments are tracked. This is
        # generally fine except in some cases (like for computing
        # init_sig) we need to produce FuncSignatures from a
        # deserialized FuncDef that lacks arguments. We won't ever
        # need to use those inside of a FuncIR, so we just make up
        # some crap.
        if hasattr(fdef, "arguments"):
            arg_names = [arg.variable.name for arg in fdef.arguments]
        else:
            arg_names = [name or "" for name in fdef.arg_names]

        args = [
            RuntimeArg(arg_name, arg_type, arg_kind, arg_pos_only)
            for arg_name, arg_kind, arg_type, arg_pos_only in zip(
                arg_names, fdef.arg_kinds, arg_types, arg_pos_onlys
            )
        ]

        # We force certain dunder methods to return objects to support letting them
        # return NotImplemented. It also avoids some pointless boxing and unboxing,
        # since tp_richcompare needs an object anyways.
        if fdef.name in ("__eq__", "__ne__", "__lt__", "__gt__", "__le__", "__ge__"):
            ret = object_rprimitive
        return FuncSignature(args, ret)

    def is_native_module(self, module: str) -> bool:
        """Is the given module one compiled by mypyc?"""
        return module in self.group_map

    def is_native_ref_expr(self, expr: RefExpr) -> bool:
        if expr.node is None:
            return False
        if "." in expr.node.fullname:
            return self.is_native_module(expr.node.fullname.rpartition(".")[0])
        return True

    def is_native_module_ref_expr(self, expr: RefExpr) -> bool:
        return self.is_native_ref_expr(expr) and expr.kind == GDEF