from typing import Dict, Optional, Union
from collections import OrderedDict

from mypy.nodes import FuncDef, TypeInfo, SymbolNode, ARG_STAR, ARG_STAR2
from mypy.types import (
    Instance, Type, CallableType, LiteralType, TypedDictType, UnboundType, PartialType,
    UninhabitedType, Overloaded, UnionType, TypeType, AnyType, NoneTyp, TupleType, TypeVarType,
    get_proper_type
)

from mypyc.ops import (
    RType, RUnion, RTuple, RInstance, object_rprimitive, dict_rprimitive, tuple_rprimitive,
    none_rprimitive, int_rprimitive, float_rprimitive, str_rprimitive, bool_rprimitive,
    list_rprimitive, set_rprimitive, FuncSignature, ClassIR, FuncDecl, RuntimeArg, LiteralsMap
)


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

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

    def __init__(self, group_map: Dict[str, Optional[str]]) -> None:
        self.group_map = group_map
        self.type_to_ir = {}  # type: Dict[TypeInfo, ClassIR]
        self.func_to_decl = {}  # type: Dict[SymbolNode, FuncDecl]
        # LiteralsMap maps literal values to a static name. Each
        # compilation group has its own LiteralsMap. (Since they can't
        # share literals.)
        self.literals = {
            v: OrderedDict() for v in group_map.values()
        }  # type: Dict[Optional[str], LiteralsMap]

    def type_to_rtype(self, typ: Optional[Type]) -> 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.str':
                return str_rprimitive
            elif typ.type.fullname == 'builtins.bool':
                return bool_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 in self.type_to_ir:
                return RInstance(self.type_to_ir[typ.type])
            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: int) -> 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)]
            ret = self.type_to_rtype(fdef.type.ret_type)
        else:
            # Handle unannotated functions
            arg_types = [object_rprimitive for arg in fdef.arguments]
            ret = object_rprimitive

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

        # 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 literal_static_name(self, module: str,
                            value: Union[int, float, complex, str, bytes]) -> str:
        # Literals are shared between modules in a compilation group
        # but not outside the group.
        literals = self.literals[self.group_map.get(module)]

        # Include type to distinguish between 1 and 1.0, and so on.
        key = (type(value), value)
        if key not in literals:
            if isinstance(value, str):
                prefix = 'unicode_'
            else:
                prefix = type(value).__name__ + '_'
            literals[key] = prefix + str(len(literals))
        return literals[key]
