import torch
from torch import Tensor
import os
import importlib.machinery
import inspect
from datetime import datetime
import sys
# sys.setrecursionlimit(10000)

def enable_gen_case():
    def check_env_flag(name, default=''):
        return os.getenv(name, default).upper() in ['ON', '1', 'YES', 'TRUE', 'Y']
    enable_gen_case = check_env_flag('TMO_GEN_CASE')
    enalbe_dump_data = check_env_flag('TMO_GEN_CASE_DUMP_DATA')
    tmo_gen_case_op_name = os.getenv('TMO_GEN_CASE_OP_NAME')
    op_names = tmo_gen_case_op_name.split(';') if tmo_gen_case_op_name else []
    tmo_gen_case_path = os.getenv('TMO_GEN_CASE_PATH', os.getcwd())
    return enable_gen_case, op_names, tmo_gen_case_path, enalbe_dump_data

ENABLE_GEN_CASE, GEN_CASE_OP_NAMES, GEN_CASE_PATH, ENALBE_DUMP_DATA = enable_gen_case()

def gen_case_dump(**kwargs):
    torch.save(kwargs['map'], kwargs['path'])

def dump_info(value):
    info = dict()
    if value.__class__ in (list, tuple):
        cls_set = {elem.__class__ for elem in value}
        if {torch.Tensor, list, tuple} & cls_set:
            temp_li = []
            for elem in value:
                temp_li.append(dump_info(elem))
            info = {'type': value.__class__, 'has_compound': True, 'data': temp_li}
        else:
            info = {'type': value.__class__, 'has_compound': False, 'data': value}
    elif value.__class__ is dict:
        temp_dic = dict()
        for k, v in value.items():
            temp_dic[k] = dump_info(v)
        info = {'type': value.__class__, 'data': temp_dic}
    elif value.__class__ is torch.Tensor:
        if value is None:
            info = {'type': value.__class__, 'data': value}
        else:
            info = {'type': value.__class__,
                    'shape': value.shape,
                    'dtype': value.dtype,
                    'device': value.device,
                    'is_contiguous': value.is_contiguous(),
                    'stride': value.stride(),
                    'data': value if value.dtype in (torch.int16, torch.int32, torch.int64) else ''}
    else:
        info = {'type': value.__class__, 'data': value}
    return info

def dump_pt_case(func, *args, **kwargs):
    op_name = func.__name__
    case_dir = os.path.join(GEN_CASE_PATH, 'tmo_gen_case', op_name)
    os.makedirs(case_dir, exist_ok=True)
    time_stamp = int(datetime.now().timestamp()*1e6)
    case_path = os.path.join(case_dir, op_name + '_' + str(time_stamp) + '.pt')
    print("[torch_mlu_ops] dump case ====> ", case_path)

    signature = inspect.signature(func)
    param_objs = [obj for obj in signature.parameters.values()]  # signature.parameters是有序的mapping
    params_map = {"op": op_name, "dump_data": ENALBE_DUMP_DATA}
    if ENALBE_DUMP_DATA:
        for obj, value in zip(param_objs[:len(args)], args):
            params_map[obj.name] = value
        for obj in param_objs[len(args):]:
            params_map[obj.name] = obj.default
        params_map = {**params_map, **kwargs}
    else:
        for obj, value in zip(param_objs[:len(args)], args):
            params_map[obj.name] = dump_info(value)
        for obj in param_objs[len(args):]:
            params_map[obj.name] = dump_info(obj.default)
        for k, v in kwargs.items():
            if k in signature.parameters.keys():
                params_map[k] = dump_info(v)
    gen_case_dump(map = params_map, path = case_path)

def dump_case(func):
    def wrapper_dump(*args, **kwargs):
        dump_pt_case(func, *args, **kwargs)
        return func(*args, **kwargs)
    def wrapper_nothing(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper_dump if (ENABLE_GEN_CASE and \
                            (GEN_CASE_OP_NAMES == [] or func.__name__ in GEN_CASE_OP_NAMES)) \
                        else wrapper_nothing

def add_gen_case_decorator(frame):
    if ENABLE_GEN_CASE:
        _current_module = inspect.getmodule(frame)
        members = dict()
        for k, v in inspect.getmembers(_current_module):
            members[k] = v

        func_names = members['__all__'] if GEN_CASE_OP_NAMES == [] else GEN_CASE_OP_NAMES
        for name in func_names:
            setattr(_current_module, name, dump_case(members[name]))

dtype2str_map = { torch.half: "half",
                  torch.float16: "half",
                  torch.float: "float",
                  torch.bfloat16: "bfloat16",
                  torch.int32: "int32",
                  torch.int8: "int8"
                }

def torchDtype2Str(torch_dtype: torch.dtype) -> str:
    assert torch_dtype in dtype2str_map, "unrecognized torch type: {}".format(torch_dtype)
    return dtype2str_map[torch_dtype]


def get_custom_op_library_path():
    loader_details = (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES)
    ext_specs = importlib.machinery.FileFinder(os.path.dirname(__file__), loader_details).find_spec("_C")
    assert ext_specs is not None, "Could not find custom op library"
    return ext_specs.origin
