from __future__ import annotations

import inspect
import warnings
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union

import torch

from .contexts.backends import EditBackend
from .contexts.Tracer import Tracer
from .intervention import InterventionProtocol, InterventionProxy
from .tracing import protocols


class Envoy:
    """Envoy object act as proxies for torch modules within a model's module tree in order to add nnsight functionality.
    Proxies of the underlying module's output and input are accessed by `.output` and `.input` respectively.

    Attributes:
        path (str): String representing the attribute path of this Envoy's module relative the the root model. Separated by '.' e.x ('transformer.h.0.mlp'). Set by NNsight on initialization of meta model.
        _fake_outputs (List[torch.Tensor]): List of 'meta' tensors built from the outputs most recent _scan. Is list as there can be multiple shapes for a module called more than once.
        _fake_inputs (List[torch.Tensor]): List of 'meta' tensors built from the inputs most recent _scan. Is list as there can be multiple shapes for a module called more than once.
        output (nnsight.intervention.InterventionProxy): Proxy object representing the output of this Envoy's module. Reset on forward pass.
        input (nnsight.intervention.InterventionProxy): Proxy object representing the input of this Envoy's module. Reset on forward pass.
        _call_iter (int): Integer representing the current iteration of this Envoy's module's inputs/outputs.
        _tracer (nnsight.context.Tracer.Tracer): Object which adds this Envoy's module's output and input proxies to an intervention graph. Must be set on Envoys objects manually by the Tracer.
    """

    def __init__(self, module: torch.nn.Module, module_path: str = ""):

        self.path = module_path

        self._fake_outputs: List[torch.Tensor] = []
        self._fake_inputs: List[torch.Tensor] = []

        self._output: Optional[InterventionProxy] = None
        self._input: Optional[InterventionProxy] = None

        self._call_iter = 0

        self._tracer: Tracer = None

        self._module = module
        self._sub_envoys: List[Envoy] = []

        # Register hook on underlying module to update the _fake_outputs and _fake_inputs on forward pass.
        self._hook_handle = self._module.register_forward_hook(
            self._hook, with_kwargs=True
        )

        for name, module in self._module.named_children():

            setattr(self, name, module)

    def _update(self, module: torch.nn.Module) -> None:
        """Updates the ._model attribute using a new model of the same architecture.
        Used when loading the real weights (dispatching) and need to replace the underlying modules.
        """

        self._module = module

        self._hook_handle.remove()

        self._hook_handle = self._module.register_forward_hook(
            self._hook, with_kwargs=True
        )

        for i, module in enumerate(self._module.children()):

            self._sub_envoys[i]._update(module)

    def _add_envoy(self, module: torch.nn.Module, name: str) -> None:
        """Adds a new Envoy for a given torch module under this Envoy.

        Args:
            module (torch.nn.Module): Module to create Envoy for.
            name (str): name of envoy/attribute.
        """

        envoy = Envoy(module, module_path=f"{self.path}.{name}")

        self._sub_envoys.append(envoy)

        # If the module already has a sub-module named 'input' or 'output',
        # mount the proxy access to 'nns_input' or 'nns_output instead.
        if hasattr(Envoy, name):

            self._handle_overloaded_mount(envoy, name)

        else:

            super().__setattr__(name, envoy)

    def _handle_overloaded_mount(self, envoy: Envoy, mount_point: str) -> None:
        """If a given module already has an attribute of the same name as something nnsight wants to add, we need to rename it.

        Directly edits the underlying class to accomplish this.

        Args:
            envoy (Envoy): Envoy to handle.
            mount_point (str): Overloaded attribute name.
        """

        warnings.warn(
            f"Module of type `{type(self._module)}` has pre-defined a `{mount_point}` attribute. nnsight access for `{mount_point}` will be mounted at `.nns_{mount_point}` instead of `.{mount_point}` for this module only."
        )

        # If we already shifted a mount point dont create another new class.
        if "Preserved" in self.__class__.__name__:

            new_cls = self.__class__

        else:

            new_cls = type(
                f"{Envoy.__name__}.Preserved",
                (Envoy,),
                {},
            )

        # Get the normal proxy mount point
        mount = getattr(new_cls, mount_point)

        # Move it to nns_<mount point>
        setattr(new_cls, f"nns_{mount_point}", mount)
        # Set the sub-module/envoy to the normal mount point on the CLASS itself not the instance.
        setattr(new_cls, mount_point, envoy)

        # Update the class on the instance
        self.__class__ = new_cls

    def _set_tracer(self, tracer: Tracer, propagate=True):
        """Set tracer object on Envoy.

        Args:
            tracer (Tracer): Tracer to set.
            propagate (bool, optional): If to propagate to all sub-modules. Defaults to True.
        """

        self._tracer = tracer

        if propagate:
            for envoy in self._sub_envoys:
                envoy._set_tracer(tracer, propagate=True)

    def _scanning(self) -> bool:
        """Whether or not in scanning mode. Checks the current Tracer's Invoker.

        Returns:
            bool: _description_
        """

        try:

            return self._tracer.invoker.scanning

        except:

            return False

    def _reset_proxies(self, propagate: bool = True) -> None:
        """Sets proxies to None.

        Args:
            propagate (bool, optional): If to propagate to all sub-modules. Defaults to True.
        """

        self._output: InterventionProxy = None
        self._input: InterventionProxy = None

        if propagate:
            for envoy in self._sub_envoys:
                envoy._reset_proxies(propagate=True)

    def _reset(self, propagate: bool = True) -> None:
        """Sets _call_iter to zero. Calls ._reset_proxies as well.

        Args:
            propagate (bool, optional): If to propagate to all sub-modules. Defaults to True.
        """

        self._reset_proxies(propagate=False)

        self._call_iter = 0

        if propagate:
            for envoy in self._sub_envoys:
                envoy._reset(propagate=True)

    def _clear(self, propagate: bool = True) -> None:
        """Clears _fake_outputs and _fake_inputs. Calls ._reset as well.

        Args:
            propagate (bool, optional): If to propagate to all sub-modules. Defaults to True.
        """

        self._reset(propagate=False)

        self._fake_outputs = []
        self._fake_inputs = []

        if propagate:
            for envoy in self._sub_envoys:
                envoy._clear(propagate=True)

    def _hook(
        self, module: torch.nn.Module, input: Any, input_kwargs: Dict, output: Any
    ):

        if self._scanning():

            self._reset_proxies(propagate=False)

            input = (input, input_kwargs)

            self._fake_outputs.append(output)
            self._fake_inputs.append(input)

    def next(self, increment: int = 1, propagate: bool = False) -> Envoy:

        self._call_iter += increment

        self._reset_proxies(propagate=False)

        if propagate:
            for envoy in self._sub_envoys:
                envoy.next(increment=increment, propagate=True)

        return self

    def to(self, *args, **kwargs) -> Envoy:
        """Override torch.nn.Module.to so this returns the Envoy, not the underlying module when doing: model = model.to(...)

        Returns:
            Envoy: Envoy.
        """

        self._module = self._module.to(*args, **kwargs)

        return self

    def modules(
        self, include_fn: Callable[[Envoy], bool] = None, names: bool = False, envoys: List = None
    ) -> List[Envoy]:
        """Returns all Envoys in the Envoy tree.

        Args:
            include_fn (Callable, optional): Optional function to be ran against all Envoys to check if they should be included in the final collection of Envoys. Defaults to None.
            names (bool, optional): If to include the name/module_path of returned Envoys along with the Envoy itself. Defaults to False.

        Returns:
            List[Envoy]: Included Envoys
        """

        if envoys is None:
            envoys = list()

        included = True

        if include_fn is not None:
            included = include_fn(self)

        if included:
            if names:
                envoys.append((self.path, self))
            else:
                envoys.append(self)

        for sub_envoy in self._sub_envoys:
            sub_envoy.modules(include_fn=include_fn, names=names, envoys=envoys)

        return envoys

    def named_modules(self, *args, **kwargs) -> List[Tuple[str, Envoy]]:
        """Returns all Envoys in the Envoy tree along with their name/module_path.

        Args:
            include_fn (Callable, optional): Optional function to be ran against all Envoys to check if they should be included in the final collection of Envoys. Defaults to None.

        Returns:
            List[Tuple[str, Envoy]]: Included Envoys and their names/module_paths.
        """

        return self.modules(*args, **kwargs, names=True)

    def _repr_module_list(self):

        list_of_reprs = [repr(item) for item in self._sub_envoys]
        if len(list_of_reprs) == 0:
            return self._module._get_name() + "()"

        start_end_indices = [[0, 0]]
        repeated_blocks = [list_of_reprs[0]]
        for i, r in enumerate(list_of_reprs[1:], 1):
            if r == repeated_blocks[-1]:
                start_end_indices[-1][1] += 1
                continue

            start_end_indices.append([i, i])
            repeated_blocks.append(r)

        lines = []
        main_str = self._module._get_name() + "("
        for (start_id, end_id), b in zip(start_end_indices, repeated_blocks):
            local_repr = f"({start_id}): {b}"  # default repr

            if start_id != end_id:
                n = end_id - start_id + 1
                local_repr = f"({start_id}-{end_id}): {n} x {b}"

            local_repr = torch.nn.modules.module._addindent(local_repr, 2)
            lines.append(local_repr)

        main_str += "\n  " + "\n  ".join(lines) + "\n"
        main_str += ")"
        return main_str

    def __repr__(self) -> str:
        """Wrapper method for underlying module's string representation.

        Returns:
            str: String.
        """

        if isinstance(self._module, torch.nn.ModuleList):

            return self._repr_module_list()

        extra_lines = []
        extra_repr = self._module.extra_repr()
        # empty string will be split into list ['']
        if extra_repr:
            extra_lines = extra_repr.split("\n")
        child_lines = []
        for attribute_name, attribute in self.__dict__.items():

            if attribute_name == "_tracer":
                continue

            if isinstance(attribute, Envoy):

                mod_str = repr(attribute)
                mod_str = torch.nn.modules.module._addindent(mod_str, 2)
                child_lines.append("(" + attribute_name + "): " + mod_str)

        lines = extra_lines + child_lines

        main_str = self._module._get_name() + "("
        if lines:
            # simple one-liner info, which most builtin Modules will use
            if len(extra_lines) == 1 and not child_lines:
                main_str += extra_lines[0]
            else:
                main_str += "\n  " + "\n  ".join(lines) + "\n"

        main_str += ")"

        return main_str

    def __iter__(self) -> Iterator[Envoy]:
        """Wrapper method for underlying ModuleList iterator.

        Returns:
            Iterator[Envoy]: Iterator.
        """

        return iter(self._sub_envoys)

    def __getitem__(self, key: int) -> Envoy:
        """Wrapper method for underlying ModuleList getitem.

        Args:
            key (int): Key.

        Returns:
            Envoy: Envoy.
        """

        return self._sub_envoys[key]

    def __len__(self) -> int:
        """Wrapper method for underlying ModuleList len.

        Returns:
            int: Length.
        """

        return len(self._module)

    def __getattr__(self, key: str) -> Union[Envoy, Any]:
        """Wrapper method for underlying module's attributes.

        Args:
            key (str): Key.

        Returns:
            Any: Attribute.
        """

        return getattr(self._module, key)

    def __setattr__(self, key: Any, value: Any) -> None:
        """Overload setattr to create and set an Envoy when trying to set a torch Module."""

        if key != "_module" and isinstance(value, torch.nn.Module):

            setattr(self._module, key, value)

            self._add_envoy(value, key)

        else:

            super().__setattr__(key, value)

    def __call__(
        self, *args: List[Any], hook=False, **kwargs: Dict[str, Any]
    ) -> InterventionProxy:
        """Creates a proxy to call the underlying module's forward method with some inputs.

        Returns:
            InterventionProxy: Module call proxy.
        """

        if isinstance(self._tracer.backend, EditBackend):
            hook = True

        return protocols.ApplyModuleProtocol.add(
            self._tracer.graph, self.path, *args, hook=hook, **kwargs
        )

    @property
    def output(self) -> InterventionProxy:
        """
        Calling denotes the user wishes to get the output of the underlying module and therefore we create a Proxy of that request.
        Only generates a proxy the first time it is references otherwise return the already set one.

        Returns:
            InterventionProxy: Output proxy.
        """
        if self._output is None:

            if isinstance(self._module, torch.nn.ModuleList):

                self._output = [envoy.output for envoy in self._sub_envoys]

                return self._output

            if len(self._fake_outputs) == 0:
                fake_output = inspect._empty
            elif self._call_iter >= len(self._fake_outputs):
                # TODO warning?
                fake_output = self._fake_outputs[-1]
            else:
                fake_output = self._fake_outputs[self._call_iter]

            module_path = f"{self.path}.output"

            self._output = InterventionProtocol.add(
                self._tracer.graph,
                fake_output,
                args=[
                    module_path,
                    len(self._tracer._invoker_inputs) - 1,
                    self._call_iter,
                ],
            )

        return self._output

    @output.setter
    def output(self, value: Union[InterventionProxy, Any]) -> None:
        """
        Calling denotes the user wishes to set the output of the underlying module and therefore we create a Proxy of that request.

        Args:
            value (Union[InterventionProxy, Any]): Value to set output to.
        """

        protocols.SwapProtocol.add(self.output.node, value)

        self._output = None

    @property
    def inputs(self) -> InterventionProxy:
        """
        Calling denotes the user wishes to get the input of the underlying module and therefore we create a Proxy of that request.
        Only generates a proxy the first time it is references otherwise return the already set one.

        Returns:
            InterventionProxy: Input proxy.
        """
        if self._input is None:

            if isinstance(self._module, torch.nn.ModuleList):

                self._input = [envoy.input for envoy in self._sub_envoys]

                return self._input

            if len(self._fake_inputs) == 0:
                fake_input = inspect._empty
            elif self._call_iter >= len(self._fake_inputs):
                # TODO warning?
                fake_input = self._fake_inputs[-1]
            else:
                fake_input = self._fake_inputs[self._call_iter]

            module_path = f"{self.path}.input"

            self._input = InterventionProtocol.add(
                self._tracer.graph,
                fake_input,
                args=[
                    module_path,
                    len(self._tracer._invoker_inputs) - 1,
                    self._call_iter,
                ],
            )

        return self._input

    @inputs.setter
    def inputs(self, value: Union[InterventionProxy, Any]) -> None:
        """
        Calling denotes the user wishes to set the input of the underlying module and therefore we create a Proxy of that request.

        Args:
            value (Union[InterventionProxy, Any]): Value to set input to.
        """

        protocols.SwapProtocol.add(self.inputs.node, value)

        self._input = None

    @property
    def input(self) -> InterventionProxy:
        """Getting the first positional argument input of the model's module.

        Returns:
            InterventionProxy: Input proxy.
        """

        return self.inputs[0][0]
    
    @input.setter
    def input(self, value: Union[InterventionProxy, Any]) -> None:
        """Setting the value of the input's first positionl argument in the model's module.
        
        Args;
            value (Union[InterventionProxy, Any]): Value to set the input to.
        """
        
        self.inputs = ((value,) + self.inputs[0][1:],) + (self.inputs[1:])