import copy
from contextlib import contextmanager

import torch
from loguru import logger


@contextmanager
def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False):
    old_register_parameter = torch.nn.Module.register_parameter
    if include_buffers:
        old_register_buffer = torch.nn.Module.register_buffer

    def register_empty_parameter(module, name, param):
        old_register_parameter(module, name, param)
        if param is not None:
            param_cls = type(module._parameters[name])
            kwargs = module._parameters[name].__dict__
            kwargs["requires_grad"] = param.requires_grad
            module._parameters[name] = param_cls(
                module._parameters[name].to(device), **kwargs
            )

    def register_empty_buffer(module, name, buffer, persistent=True):
        old_register_buffer(module, name, buffer, persistent=persistent)
        if buffer is not None:
            module._buffers[name] = module._buffers[name].to(device)

    def patch_tensor_constructor(fn):
        def wrapper(*args, **kwargs):
            kwargs["device"] = device
            return fn(*args, **kwargs)

        return wrapper

    if include_buffers:
        tensor_constructors_to_patch = {
            torch_function_name: getattr(torch, torch_function_name)
            for torch_function_name in ["empty", "zeros", "ones", "full"]
        }
    else:
        tensor_constructors_to_patch = {}

    try:
        torch.nn.Module.register_parameter = register_empty_parameter
        if include_buffers:
            torch.nn.Module.register_buffer = register_empty_buffer
        for torch_function_name in tensor_constructors_to_patch.keys():
            setattr(
                torch,
                torch_function_name,
                patch_tensor_constructor(getattr(torch, torch_function_name)),
            )
        yield
    finally:
        torch.nn.Module.register_parameter = old_register_parameter
        if include_buffers:
            torch.nn.Module.register_buffer = old_register_buffer
        for (
            torch_function_name,
            old_torch_function,
        ) in tensor_constructors_to_patch.items():
            setattr(torch, torch_function_name, old_torch_function)


def cast_to(
    weight: torch.Tensor, dtype: torch.dtype, device: torch.device
) -> torch.Tensor:
    if hasattr(weight, "__class__") and "torchao.dtypes" in str(weight.__class__):
        r = weight.to(device=device)
    elif hasattr(weight, "__class__") and "torchao.quantization" in str(
        weight.__class__
    ):
        r = weight.to(dtype=dtype, device=device)
    else:
        r = torch.empty_like(weight, dtype=dtype, device=device)
        r.copy_(weight)
    return r


def cast_to_device(weight: torch.Tensor, device: torch.device) -> torch.Tensor:
    if hasattr(weight, "__class__") and "optimum.quanto" in str(weight.__class__):
        r = weight.to(device)
        return r
    else:
        r = torch.empty_like(weight, device=device)
        r.copy_(weight)
        return r


class AutoWrappedModule(torch.nn.Module):
    def __init__(
        self,
        module: torch.nn.Module,
        offload_dtype,
        offload_device,
        onload_dtype,
        onload_device,
        computation_dtype,
        computation_device,
    ):
        super().__init__()
        self.module = module.to(dtype=offload_dtype, device=offload_device)
        self.offload_dtype = offload_dtype
        self.offload_device = offload_device
        self.onload_dtype = onload_dtype
        self.onload_device = onload_device
        self.computation_dtype = computation_dtype
        self.computation_device = computation_device
        self.state = 0

    def offload(self):
        if self.state == 1 and (
            self.offload_dtype != self.onload_dtype
            or self.offload_device != self.onload_device
        ):
            self.module.to(dtype=self.offload_dtype, device=self.offload_device)
            self.state = 0

    def onload(self):
        if self.state == 0 and (
            self.offload_dtype != self.onload_dtype
            or self.offload_device != self.onload_device
        ):
            self.module.to(dtype=self.onload_dtype, device=self.onload_device)
            self.state = 1

    def forward(self, *args, **kwargs):
        if (
            self.onload_dtype == self.computation_dtype
            and self.onload_device == self.computation_device
        ):
            module = self.module
        else:
            module = copy.deepcopy(self.module).to(
                dtype=self.computation_dtype, device=self.computation_device
            )
        r = module(*args, **kwargs)
        return r

    def __getattr__(self, name):
        # Proxy attribute access to wrapped module for non-special attributes
        if name in [
            "module",
            "offload_dtype",
            "offload_device",
            "onload_dtype",
            "onload_device",
            "computation_dtype",
            "computation_device",
            "state",
        ]:
            return super().__getattr__(name)
        return getattr(self.module, name)


class AutoWrappedLinear(torch.nn.Linear):
    def __init__(
        self,
        module: torch.nn.Linear,
        offload_dtype,
        offload_device,
        onload_dtype,
        onload_device,
        computation_dtype,
        computation_device,
    ):
        with init_weights_on_device(device=torch.device("meta")):
            super().__init__(
                in_features=module.in_features,
                out_features=module.out_features,
                bias=module.bias is not None,
                dtype=offload_dtype,
                device=offload_device,
            )
        self.weight = module.weight
        self.bias = module.bias
        self.offload_dtype = offload_dtype
        self.offload_device = offload_device
        self.onload_dtype = onload_dtype
        self.onload_device = onload_device
        self.computation_dtype = computation_dtype
        self.computation_device = computation_device
        self.state = 0

    def offload(self):
        if self.state == 1 and (
            self.offload_dtype != self.onload_dtype
            or self.offload_device != self.onload_device
        ):
            self.to(dtype=self.offload_dtype, device=self.offload_device)
            self.state = 0

    def onload(self):
        if self.state == 0 and (
            self.offload_dtype != self.onload_dtype
            or self.offload_device != self.onload_device
        ):
            self.to(dtype=self.onload_dtype, device=self.onload_device)
            self.state = 1

    def forward(self, x, *args, **kwargs):
        if (
            self.onload_dtype == self.computation_dtype
            and self.onload_device == self.computation_device
        ):
            weight, bias = self.weight, self.bias
        else:
            weight = cast_to(
                self.weight, self.computation_dtype, self.computation_device
            )
            bias = (
                None
                if self.bias is None
                else cast_to(self.bias, self.computation_dtype, self.computation_device)
            )
        r = torch.nn.functional.linear(x, weight, bias)
        return r


def enable_vram_management_recursively(
    model: torch.nn.Module,
    module_map: dict,
    module_config: dict,
    max_num_param=None,
    overflow_module_config: dict = None,
    total_num_param=0,
):
    for name, module in model.named_children():
        module_type = module.__class__.__name__
        for source_module, target_module in module_map.items():
            if isinstance(module, source_module):
                num_param = sum(p.numel() for p in module.parameters())
                is_overflow = (
                    max_num_param is not None
                    and total_num_param + num_param > max_num_param
                )
                if is_overflow:
                    module_config_ = overflow_module_config
                else:
                    module_config_ = module_config
                module_ = target_module(module, **module_config_)
                setattr(model, name, module_)
                total_num_param += num_param
                break
        else:
            total_num_param = enable_vram_management_recursively(
                module,
                module_map,
                module_config,
                max_num_param,
                overflow_module_config,
                total_num_param,
            )
    return total_num_param


def enable_vram_management(
    model: torch.nn.Module,
    module_map: dict,
    module_config: dict,
    max_num_param=None,
    overflow_module_config: dict = None,
):
    total_num_param = enable_vram_management_recursively(
        model,
        module_map,
        module_config,
        max_num_param,
        overflow_module_config,
        total_num_param=0,
    )
    model.vram_management_enabled = True
    max_persistent_str = f"{max_num_param / 1e6:.2f}M" if max_num_param else "None"
    logger.info(
        f"[VRAM] Total wrapped parameters: {total_num_param / 1e6:.2f}M "
        f"(max_persistent: {max_persistent_str})"
    )
