import inspect
import time
import types
import warnings
from functools import wraps
from typing import Callable, Dict, List, TYPE_CHECKING, Tuple

import torch

from .constants import ORIG_TORCH_FUNCS
from .helper_funcs import (
    get_vars_of_type_from_obj,
    identity,
    log_current_rng_states,
    make_random_barcode,
    nested_getattr,
    print_override,
    safe_copy,
)
from .logging_funcs import log_function_output_tensors, log_source_tensor

if TYPE_CHECKING:
    from model_history import ModelHistory

funcs_not_to_log = ["numpy", "__array__", "size", "dim"]
print_funcs = ["__repr__", "__str__", "_str"]


def torch_func_decorator(self, func: Callable, func_name: str):
    @wraps(func)
    def wrapped_func(*args, **kwargs):
        # Initial bookkeeping; check if it's a special function, organize the arguments.
        self.current_function_call_barcode = 0
        if (func_name in funcs_not_to_log) or (not self._track_tensors) or self._pause_logging:
            out = func(*args, **kwargs)
            return out
        all_args = list(args) + list(kwargs.values())
        arg_tensorlike = get_vars_of_type_from_obj(all_args, torch.Tensor)

        # Register any buffer tensors in the arguments.

        for t in arg_tensorlike:
            if hasattr(t, "tl_buffer_address"):
                log_source_tensor(self, t, "buffer", getattr(t, "tl_buffer_address"))

        if (func_name in print_funcs) and (len(arg_tensorlike) > 0):
            out = print_override(args[0], func_name)
            return out

        # Copy the args and kwargs in case they change in-place:
        if self.save_function_args:
            arg_copies = tuple([safe_copy(arg) for arg in args])
            kwarg_copies = {k: safe_copy(v) for k, v in kwargs.items()}
        else:
            arg_copies = args
            kwarg_copies = kwargs

        # Call the function, tracking the timing, rng states, and whether it's a nested function
        func_call_barcode = make_random_barcode()
        self.current_function_call_barcode = func_call_barcode
        start_time = time.time()
        func_rng_states = log_current_rng_states()
        out_orig = func(*args, **kwargs)
        func_time_elapsed = time.time() - start_time
        is_bottom_level_func = self.current_function_call_barcode == func_call_barcode

        if func_name in ["__setitem__", "zero_", "__delitem__"]:
            out_orig = args[0]

        was_inplace = len(args) > 0 and id(out_orig) == id(args[0])
        if was_inplace:
            out_orig = safe_copy(out_orig)

        # Log all output tensors
        output_tensors = get_vars_of_type_from_obj(
            out_orig,
            which_type=torch.Tensor,
            subclass_exceptions=[torch.nn.Parameter],
        )

        if len(output_tensors) > 0:
            log_function_output_tensors(
                self,
                func,
                func_name,
                args,
                kwargs,
                arg_copies,
                kwarg_copies,
                out_orig,
                func_time_elapsed,
                func_rng_states,
                is_bottom_level_func,
            )

            # Propagate label to original tensor for in-place ops so
            # subsequent operations see the updated label.
            if was_inplace:
                if hasattr(out_orig, "tl_tensor_label_raw"):
                    args[0].tl_tensor_label_raw = out_orig.tl_tensor_label_raw

        return out_orig

    return wrapped_func


def decorate_pytorch(
    self: "ModelHistory", torch_module: types.ModuleType, orig_func_defs: List[Tuple]
) -> Dict[Callable, Callable]:
    """Mutates all PyTorch functions (TEMPORARILY!) to save the outputs of any functions
    that return Tensors, along with marking them with metadata. Returns a list of tuples that
    save the current state of the functions, such that they can be restored when done.

    args:
        torch_module: The top-level torch module (i.e., from "import torch").
            This is supplied as an argument on the off-chance that the user has imported torch
            and done their own monkey-patching.
        tensors_to_mutate: A list of tensors that will be mutated (since any tensors created
            before calling the torch mutation function will not be mutated).
        orig_func_defs: Supply a list from outside to guarantee it can be cleaned up properly.
        tensor_record: A list to which the outputs of the functions will be appended.

    returns:
        List of tuples consisting of [namespace, func_name, orig_func], sufficient
        to return torch to normal when finished, and also a dict mapping mutated functions to original functions.
    """

    # Do a pass to save the original func defs.
    collect_orig_func_defs(torch_module, orig_func_defs)
    decorated_func_mapper = {}

    # Get references to the function classes.
    function_class = type(lambda: 0)
    builtin_class = type(torch.mean)
    method_class = type(torch.Tensor.__add__)
    wrapper_class = type(torch.Tensor.__getitem__)
    getset_class = type(torch.Tensor.real)

    for namespace_name, func_name in ORIG_TORCH_FUNCS:
        namespace_name_notorch = namespace_name.replace("torch.", "")
        local_func_namespace = nested_getattr(torch_module, namespace_name_notorch)
        if not hasattr(local_func_namespace, func_name):
            continue
        orig_func = getattr(local_func_namespace, func_name)
        if func_name not in self.func_argnames:
            get_func_argnames(self, orig_func, func_name)
        if getattr(orig_func, "__name__", False) == "wrapped_func":
            continue

        if type(orig_func) in [function_class, builtin_class, method_class, wrapper_class]:
            new_func = torch_func_decorator(self, orig_func, func_name)
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    setattr(local_func_namespace, func_name, new_func)
                new_func.tl_is_decorated_function = True
                decorated_func_mapper[new_func] = orig_func
                decorated_func_mapper[orig_func] = new_func
            except (AttributeError, TypeError):
                pass

        elif type(orig_func) is getset_class:
            getter_orig, setter_orig, deleter_orig = (
                orig_func.__get__,
                orig_func.__set__,
                orig_func.__delete__,
            )
            getter_dec, setter_dec, deleter_dec = (
                torch_func_decorator(self, getter_orig, func_name),
                torch_func_decorator(self, setter_orig, func_name),
                torch_func_decorator(self, deleter_orig, func_name),
            )
            getter_dec.tl_is_decorated_function = True
            setter_dec.tl_is_decorated_function = True
            deleter_dec.tl_is_decorated_function = True
            new_property = property(getter_dec, setter_dec, deleter_dec, doc=func_name)
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    setattr(local_func_namespace, func_name, new_property)
            except (AttributeError, TypeError) as _:
                pass
            decorated_func_mapper[new_property] = orig_func
            decorated_func_mapper[orig_func] = new_property

    # Bolt on the identity function
    new_identity = torch_func_decorator(self, identity, "identity")
    torch.identity = new_identity

    return decorated_func_mapper


def undecorate_pytorch(
    torch_module, orig_func_defs: List[Tuple], input_tensors: List[torch.Tensor]
):
    """
    Returns all PyTorch functions back to the definitions they had when mutate_pytorch was called.
    This is done for the output tensors and history_dict too to avoid ugliness. Also deletes
    the mutant versions of the functions to remove any references to old ModelHistory object.

    args:
        torch_module: The torch module object.
        orig_func_defs: List of tuples consisting of [namespace_name, func_name, orig_func], sufficient
            to regenerate the original functions.
        input_tensors: List of input tensors whose fucntions will be undecorated.
        decorated_func_mapper: Maps the decorated function to the original function
    """
    for namespace_name, func_name, orig_func in orig_func_defs:
        namespace_name_notorch = namespace_name.replace("torch.", "")
        local_func_namespace = nested_getattr(torch_module, namespace_name_notorch)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            decorated_func = getattr(local_func_namespace, func_name)
        del decorated_func
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                setattr(local_func_namespace, func_name, orig_func)
        except (AttributeError, TypeError) as _:
            continue
    if hasattr(torch, "identity"):
        delattr(torch, "identity")
    for input_tensor in input_tensors:
        if hasattr(input_tensor, "tl_tensor_label_raw"):
            delattr(input_tensor, "tl_tensor_label_raw")


def collect_orig_func_defs(torch_module: types.ModuleType, orig_func_defs: List[Tuple]):
    """Collects the original torch function definitions, so they can be restored after the logging is done.

    Args:
        torch_module: The top-level torch module
        orig_func_defs: List of tuples keeping track of the original function definitions
    """
    for namespace_name, func_name in ORIG_TORCH_FUNCS:
        namespace_name_notorch = namespace_name.replace("torch.", "")
        local_func_namespace = nested_getattr(torch_module, namespace_name_notorch)
        if not hasattr(local_func_namespace, func_name):
            continue
        orig_func = getattr(local_func_namespace, func_name)
        orig_func_defs.append((namespace_name, func_name, orig_func))


# TODO: hard-code some of the arg names; for example truediv, getitem, etc. Can crawl through and see what isn't working
def get_func_argnames(self, orig_func: Callable, func_name: str):
    """Attempts to get the argument names for a function, first by checking the signature, then
    by checking the documentation. Adds these names to func_argnames if it can find them,
    doesn't do anything if it can't."""
    if func_name in ["real", "imag", "T", "mT", "data", "H"]:
        return

    try:
        argnames = list(inspect.signature(orig_func).parameters.keys())
        argnames = tuple([arg.replace("*", "") for arg in argnames if arg not in ["cls", "self"]])
        self.func_argnames[func_name] = argnames
        return
    except ValueError:
        pass

    docstring = orig_func.__doc__
    if (type(docstring) is not str) or (len(docstring) == 0):  # if docstring missing, skip it
        return

    open_ind, close_ind = docstring.find("("), docstring.find(")")
    argstring = docstring[open_ind + 1 : close_ind]
    arg_list = argstring.split(",")
    arg_list = [arg.strip(" ") for arg in arg_list]
    argnames = []
    for arg in arg_list:
        argname = arg.split("=")[0]
        if argname in ["*", "/", "//", ""]:
            continue
        argname = argname.replace("*", "")
        argnames.append(argname)
    argnames = tuple([arg for arg in argnames if arg not in ["self", "cls"]])
    self.func_argnames[func_name] = argnames
    return
