# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium
import copy
from typing import Dict, List, Optional, Set, Tuple, Union

import torch
import transformers
from accelerate.utils import has_offloaded_params
from torch import nn

from ..quantization.rotation.hadamard_utils import apply_online_hadamard
from ..utils.device_telemetry import emit_device_telemetry
from ..utils.logger import setup_logger


log = setup_logger()


class StopForward(Exception):
    """Signal an intentional early stop of the forward pass."""
    pass

STOP_FORWARD_EXCEPTION = StopForward("Forwarding stopped")


def _restore_output_device(output: torch.Tensor, original_device: torch.device) -> torch.Tensor:
    if output.device == original_device:
        return output
    if original_device.type == "meta":
        # Meta tensors are placeholders with no backing storage; never move real
        # outputs back to meta during hook replay.
        return output
    return output.to(device=original_device)


def _materialize_if_meta_weight(module: nn.Module, *, input_device: torch.device) -> None:
    weight = getattr(module, "weight", None)
    if weight is None:
        return
    if getattr(weight, "is_meta", False) or weight.device.type == "meta":
        if has_offloaded_params(module):
            from ..utils.offload import undo_offload_to_disk

            restore_device = input_device if input_device.type != "meta" else torch.device("cpu")
            undo_offload_to_disk(module, device=restore_device, include_buffers=False)


# Models using conv1d: gpt2
class HookedConv1D(transformers.Conv1D):
    def __init__(self, nf: int, nx: int) -> None:
        torch.nn.Module.__init__(self)
        self.nf = nf
        self.nx = nx
        self.forward_hook = None
        self.forward_hook_last = False

    @staticmethod
    def from_conv1d(m: transformers.Conv1D):
        assert isinstance(m, transformers.Conv1D)
        custom = HookedConv1D(m.nf, m.nx)
        custom.weight = m.weight
        custom.bias = m.bias
        return custom

    @torch.inference_mode()
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        original_device = input.device
        _materialize_if_meta_weight(self, input_device=original_device)
        target_device = self.weight.data.device
        if original_device != target_device:
            input = input.to(device=target_device)
        output = super().forward(input)

        if self.forward_hook:
            self.forward_hook(self, (input,), output)
            if self.forward_hook_last:
                raise STOP_FORWARD_EXCEPTION.with_traceback(None)

        return _restore_output_device(output, original_device)

class HookedConv1d(torch.nn.Conv1d):
    def __init__(
        self,
    ) -> None:
        # in_channels,
        # out_channels,
        # kernel_size,
        # stride,
        # padding,
        # dilation,
        # groups,
        # padding_mode,
        torch.nn.Module.__init__(self)
        # TODO: call super().__init__() is too slow, need to find a better way
        # super().__init__(
        #     in_channels,
        #     out_channels,
        #     kernel_size,
        #     stride,
        #     padding,
        #     dilation,
        #     groups,
        #     padding_mode,
        # )
        self.forward_hook = None
        self.forward_hook_last = False

    @staticmethod
    def from_conv1d(m: torch.nn.Conv1d):
        assert isinstance(m, torch.nn.Conv1d)

        custom = HookedConv1d()
        custom.in_channels = m.in_channels
        custom.out_channels = m.out_channels
        custom.kernel_size = m.kernel_size
        custom.stride = m.stride
        custom.padding = m.padding
        custom.dilation = m.dilation
        custom.transposed = m.transposed
        custom.output_padding = m.output_padding
        custom.groups = m.groups
        custom.padding_mode = m.padding_mode

        custom.weight = m.weight
        custom.bias = m.bias
        return custom

    @torch.inference_mode()
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        original_device = input.device
        _materialize_if_meta_weight(self, input_device=original_device)
        target_device = self.weight.data.device
        if original_device != target_device:
            input = input.to(device=target_device)
        output = super().forward(input)
        if self.forward_hook:
            self.forward_hook(self, (input,), output)
            if self.forward_hook_last:
                raise STOP_FORWARD_EXCEPTION.with_traceback(None)
        return _restore_output_device(output, original_device)

# Models using conv2d: ovis
class HookedConv2d(torch.nn.Conv2d):
    def __init__(
        self,
    ) -> None:
        # in_channels,
        # out_channels,
        # kernel_size,
        # stride,
        # padding,
        # dilation,
        # groups,
        # padding_mode,
        torch.nn.Module.__init__(self)
        # TODO: call super().__init__() is too slow, need to find a better way
        # super().__init__(
        #     in_channels,
        #     out_channels,
        #     kernel_size,
        #     stride,
        #     padding,
        #     dilation,
        #     groups,
        #     padding_mode,
        # )
        self.forward_hook = None
        self.forward_hook_last = False

    @staticmethod
    def from_conv2d(m: torch.nn.Conv2d):
        assert isinstance(m, torch.nn.Conv2d)

        custom = HookedConv2d()
        custom.in_channels = m.in_channels
        custom.out_channels = m.out_channels
        custom.kernel_size = m.kernel_size
        custom.stride = m.stride
        custom.padding = m.padding
        custom.dilation = m.dilation
        custom.transposed = m.transposed
        custom.output_padding = m.output_padding
        custom.groups = m.groups
        custom.padding_mode = m.padding_mode

        custom.weight = m.weight
        custom.bias = m.bias
        return custom

    @torch.inference_mode()
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        original_device = input.device
        _materialize_if_meta_weight(self, input_device=original_device)
        target_device = self.weight.data.device
        if original_device != target_device:
            input = input.to(device=target_device)
        output = super().forward(input)
        if self.forward_hook:
            self.forward_hook(self, (input,), output)
            if self.forward_hook_last:
                raise STOP_FORWARD_EXCEPTION.with_traceback(None)
        return _restore_output_device(output, original_device)

# Models using transformers.conv1d: gpt2
class HookedTransformerConv1D(transformers.Conv1D):
    def __init__(self, nf: int, nx: int) -> None:
        torch.nn.Module.__init__(self)
        self.nf = nf
        self.nx = nx
        self.forward_hook = None
        self.forward_hook_last = False

    @staticmethod
    def from_conv1d(conv1d: transformers.Conv1D):
        custom = HookedConv1D(conv1d.nf, conv1d.nx)
        custom.weight = conv1d.weight
        custom.bias = conv1d.bias
        return custom

    @torch.inference_mode()
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        original_device = input.device
        _materialize_if_meta_weight(self, input_device=original_device)
        target_device = self.weight.data.device
        if original_device != target_device:
            input = input.to(device=target_device)
        output = super().forward(input)
        if self.forward_hook:
            self.forward_hook(self, (input,), output)
            if self.forward_hook_last:
                raise STOP_FORWARD_EXCEPTION.with_traceback(None)
        return _restore_output_device(output, original_device)

class HookedLinear(torch.nn.Linear):
    def __init__(self, in_features: int, out_features: int) -> None:
        # avoid calling super().__init__() as it would allocate memory based on in/out features
        torch.nn.Module.__init__(self)
        self.in_features = in_features
        self.out_features = out_features

        self.forward_hook = None
        self.forward_hook_last = False
        self.module_name = None

    @staticmethod
    def from_linear(linear: torch.nn.Linear):
        custom_linear = HookedLinear(linear.in_features, linear.out_features)
        custom_linear.weight = linear.weight
        custom_linear.bias = linear.bias
        custom_linear.online_full_had = getattr(linear, "online_full_had", False)
        custom_linear.online_partial_had = getattr(linear, "online_partial_had", False)
        custom_linear.had_dim = getattr(linear, "had_dim", -1)
        custom_linear.had_K = getattr(linear, "had_K", None)
        custom_linear.K = getattr(linear, "K", 1)
        return custom_linear

    @torch.inference_mode()
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        original_device = input.device
        _materialize_if_meta_weight(self, input_device=original_device)
        target_device = self.weight.data.device
        module_name = getattr(self, "module_name", None) or getattr(self, "full_name", None) or getattr(self, "name", None) or "unknown"
        emit_device_telemetry(
            "hooked_linear_forward",
            module=module_name,
            weight_device=target_device,
            input_device=original_device,
        )
        if original_device != target_device:
            input = input.to(device=target_device)
        input = apply_online_hadamard(
            input,
            online_full_had=getattr(self, "online_full_had", False),
            online_partial_had=getattr(self, "online_partial_had", False),
            had_K=getattr(self, "had_K", None),
            K=getattr(self, "K", 1),
            had_dim=getattr(self, "had_dim", -1),
        )
        output = super().forward(input)
        if self.forward_hook:
            self.forward_hook(self, (input,), output)
            if self.forward_hook_last:
                raise STOP_FORWARD_EXCEPTION.with_traceback(None)
        return _restore_output_device(output, original_device)


def _replace_module(module, child, name, level: int = 0, debug: bool = False) -> bool:
    level_indent = "---" * level
    instance_type = type(child)
    if debug:
        log.info(f"{level_indent} Hook: {instance_type.__name__}: {name}")

    # Replace nn.Linear with HookedLinear, except PhimoeTopKRouter which returns a tuple in forward()
    if isinstance(child, torch.nn.Linear) and child.__class__.__name__ != "PhimoeTopKRouter":
        setattr(module, name, HookedLinear.from_linear(child))
    elif isinstance(child, transformers.Conv1D):
        setattr(module, name, HookedTransformerConv1D.from_conv1d(child))
    elif isinstance(child, nn.Conv1d):
        setattr(module, name, HookedConv1d.from_conv1d(child))
    elif isinstance(child, nn.Conv2d):
        setattr(module, name, HookedConv2d.from_conv2d(child))
    else:
        if debug:
            log.error(f"{level_indent} Hook: execute_replace but layer skipped due to type not supported: {name}")
        return False

    return True


def replace_module_with_hooked_legacy(
    module,
    level: int = 0,
    quant_lm_head: bool = False,
    skip_module_paths: Optional[Set[str]] = None,
    _prefix: str = "",
):
    # if level == 0:
    #     log.info("Hooked Modules: Using legacy based config for targeting of modules")

    for name, child in module.named_children():
        if not quant_lm_head and hasattr(module, "get_output_embeddings") and child == module.get_output_embeddings():
            continue

        child_path = f"{_prefix}.{name}" if _prefix else name
        if skip_module_paths and child_path in skip_module_paths:
            continue

        if not _replace_module(module, child, name, level, quant_lm_head):
            replace_module_with_hooked_legacy(
                child,
                level=level + 1,
                quant_lm_head=quant_lm_head,
                skip_module_paths=skip_module_paths,
                _prefix=child_path,
            )

# deprecated features
def replace_module_with_hooked_tree(module, tree: Union[List,Dict] = [], level: int = 0, debug: bool = False):
    if level == 0:
        log.info("Hooked Modules: Using tree based config for accurate targeting of modules")

    tree = copy.copy(tree) # defensive copy

    # tuple represents targeted modules
    execute_replace = isinstance(tree, Tuple)
    # level indent
    level_indent = "---" * level

    for name, child in module.named_children():
        if debug:
            log.info(f"{level_indent} child name: {name}")
        if execute_replace:
            if name in tree:
                # do replace if name in tree and tree is a tuple
                _replace_module(module, child, name, level, debug)
            else:
                if debug:
                    log.warn(f"{level_indent} Hook: execute_replace but layer skipped due to not targeted: {name}")
        else:
            if isinstance(tree, Dict):
                if name in tree:
                    if isinstance(tree[name], Tuple) and name in tree[name]:
                        # do replace if name in tree and tree is a tuple
                        _replace_module(module, child, name, level, debug)
                    elif isinstance(tree[name], (Dict, Tuple, List)):
                        # follow tree node if name in tree and tree is a dict
                        if debug:
                            log.info(f"{level_indent} Hook: follow tree node: {name} -> nest into {name}")

                        replace_module_with_hooked_tree(child, tree=tree[name],
                                                        level=level+1, debug=debug)
                    else:
                        if debug:
                            log.warn(f"{level_indent} Hook: skipped unknown tree node dict: {name} vs tree: {tree}")

                elif "#" in tree and name.isdigit():
                    if debug:
                        log.info(f"{level_indent} Hook: follow tree node: {name} -> nest into {name}")
                    replace_module_with_hooked_tree(child, tree=tree["#"],
                                                    level=level+1, debug=debug)
                else:
                    if debug:
                        log.warn(f"{level_indent} Hook: skipped unknown tree node dict: {name} vs tree: {tree}")
            elif isinstance(tree, List):
                next_node = tree[0]
                if isinstance(next_node, Dict):
                    if name in next_node:
                        if debug:
                            log.info(f"{level_indent} Hook: follow tree node: {name} -> nest into {name}")
                        replace_module_with_hooked_tree(child, tree=next_node[name],
                                                        level=level + 1, debug=debug)
                    elif name.isdigit and "#" in next_node:
                        if debug:
                            log.info(f"{level_indent} Hook: follow tree node: {name} -> nest into {name}")
                        replace_module_with_hooked_tree(child, tree=next_node["#"],
                                                        level=level + 1, debug=debug)
                    else:
                        if debug:
                            log.warn(
                            f"{level_indent} Hook: skipped unknown tree node dict: {name} vs next_node: {next_node}")
                elif name == next_node or (next_node == "#" and name.isdigit()):
                    if debug:
                        log.info(f"{level_indent} Hook: follow tree node: {name} -> nest into {name}")
                    replace_module_with_hooked_tree(child, tree=tree[1:],
                                                    level=level + 1, debug=debug)
                else:
                    if debug:
                        log.warn(f"{level_indent} Hook: skipped unknown tree node list: {name} vs next_node: {next_node}")

            # if len(tree) > 0:
            #     # list or dict
            #     node = next(tree_iter)
            #     log.info(f"{level_indent} next node: {node}")
            #
            #     # matched user designed tree node/path
            #     if name == node or (node == "#" and name.isdigit()):
            #         log.info(f"{level_indent} Hook: follow tree node: {node} -> nest into {name}")
            #         replace_module_with_hooked(child, tree=tree[1:] if isinstance(tree, List) else tree[name], level=level + 1, debug=debug)
            #     # # list: simple node
            #     # elif isinstance(node, List):
            #     #     if name in node:
            #     #         log.info(f"{level_indent} Hook: follow tree node list: {node} -> {name} -> nest into {name}")
            #     #         replace_module_with_hooked(child, tree=tree[1:], level=level + 1, debug=debug)
            #     #     elif name.isdigit() and "#" in node:
            #     #         log.info(f"{level_indent} Hook: follow tree node list: {node} -> # -> nest into {name}")
            #     #         replace_module_with_hooked(child, tree=tree[1:], level=level + 1, debug=debug)
            #     #     else:
            #     #         log.warn(f"{level_indent} Hook: stopped at unknown tree node list: {node} -> {name}")
            #     # # dict: nested node with final modules larget as value [list]
            #     # elif isinstance(node, Dict):
            #     #     if name in node:
            #     #         log.info(f"{level_indent} Hook: follow tree node dict: {node} -> {name} -> nest into {name}")
            #     #         replace_module_with_hooked(child, tree=node[name], level=level + 1, debug=debug)
            #     #     elif name.isdigit() and "#" in node:
            #     #         log.info(f"{level_indent} Hook: follow tree node dict: {node} -> # -> nest into {name}")
            #     #         replace_module_with_hooked(child, tree=node["#"], level=level + 1, debug=debug)
            #     #     else:
            #     #         log.warn(f"{level_indent} Hook: stopped at unknown tree node dict: {node} -> {name}")
            #     else:
            #         log.warn(f"{level_indent} Hook: skipped at unknown tree node: {name}")
            # else:
            #     log.warn(f"{level_indent} Hook: follow naively -> nest into {name}")
            #     replace_module_with_hooked(child, [], level=level + 1, debug=debug)