from collections.abc import Generator
from dataclasses import dataclass

import torch
import torch.nn as nn

from vllm.logger import init_logger
from vllm.model_executor.offloader.base import BaseOffloader
from vllm_ascend.model_executor.offloader.selection import select_offload_layers

logger = init_logger(__name__)


@dataclass
class _ParamState:
    name: str
    param: nn.Parameter
    cpu_tensor: torch.Tensor
    original_device: torch.device
    num_bytes: int


class _SyncModuleOffloader:
    def __init__(self, *, module: nn.Module, layer_idx: int, offload_params: set[str]):
        self.module = module
        self.layer_idx = layer_idx
        self.param_states: list[_ParamState] = []
        self.offload_params = offload_params
        self.offloaded_bytes = 0
        self._original_forward = module.forward
        self._capture_params()
        self._install_forward_hook()

    def _matches_param(self, name: str) -> bool:
        if not self.offload_params:
            return True
        return any(f".{part}." in f".{name}." for part in self.offload_params)

    @staticmethod
    def _tensor_nbytes(tensor: torch.Tensor) -> int:
        return tensor.numel() * tensor.element_size()

    def _capture_params(self) -> None:
        for name, param in self.module.named_parameters(recurse=True):
            if not self._matches_param(name):
                continue
            cpu_tensor = param.detach().cpu().contiguous()
            self.param_states.append(
                _ParamState(
                    name=name,
                    param=param,
                    cpu_tensor=cpu_tensor,
                    original_device=param.device,
                    num_bytes=self._tensor_nbytes(cpu_tensor),
                )
            )
            param.data = cpu_tensor

    def post_init(self) -> None:
        self.offloaded_bytes = 0
        for state in self.param_states:
            state.cpu_tensor = state.param.detach().cpu().contiguous()
            state.num_bytes = self._tensor_nbytes(state.cpu_tensor)
            state.param.data = state.cpu_tensor
            self.offloaded_bytes += state.num_bytes

    @staticmethod
    def _device_from_args(args, kwargs) -> torch.device | None:
        tensors = list(args) + list(kwargs.values())
        for value in tensors:
            if isinstance(value, torch.Tensor) and value.device.type != "cpu":
                return value.device
            if isinstance(value, (tuple, list)):
                for item in value:
                    if isinstance(item, torch.Tensor) and item.device.type != "cpu":
                        return item.device
        return None

    def _onload(self, args, kwargs) -> None:
        runtime_device = self._device_from_args(args, kwargs)
        for state in self.param_states:
            target_device = state.original_device
            if target_device.type == "cpu" and runtime_device is not None:
                target_device = runtime_device
            if target_device.type == "cpu":
                continue
            if state.param.device != target_device:
                state.param.data = state.cpu_tensor.to(
                    device=target_device,
                    non_blocking=False,
                )

    def _offload(self) -> None:
        for state in self.param_states:
            if state.param.device.type != "cpu":
                state.param.data = state.cpu_tensor

    def _install_forward_hook(self) -> None:
        def forward(*args, **kwargs):
            self._onload(args, kwargs)
            try:
                return self._original_forward(*args, **kwargs)
            finally:
                self._offload()

        self.module.forward = forward


class AscendSyncOffloader(BaseOffloader):
    def __init__(
        self,
        *,
        group_size: int,
        num_in_group: int,
        offload_params: set[str] | None = None,
    ):
        self.group_size = group_size
        self.num_in_group = num_in_group
        self.offload_params = offload_params or set()
        self.selected_layer_indices: tuple[int, ...] = ()
        self.module_offloaders: list[_SyncModuleOffloader] = []
        self.total_offloaded_bytes = 0

    def wrap_modules(
        self,
        modules_generator: Generator[nn.Module, None, None],
    ) -> list[nn.Module]:
        modules = list(modules_generator)
        selection = select_offload_layers(
            num_layers=len(modules),
            group_size=self.group_size,
            num_in_group=self.num_in_group,
        )
        self.selected_layer_indices = selection.layer_indices
        if not selection.enabled:
            return modules

        for layer_idx in selection.layer_indices:
            module = modules[layer_idx]
            if not any(True for _ in module.parameters(recurse=True)):
                continue
            module_offloader = _SyncModuleOffloader(
                module=module,
                layer_idx=layer_idx,
                offload_params=self.offload_params,
            )
            if not module_offloader.param_states:
                continue
            self.module_offloaders.append(module_offloader)

        logger.info(
            "[AscendSyncOffloader] selected_layers=%s, wrapped=%d, "
            "group_size=%d, num_in_group=%d",
            self.selected_layer_indices,
            len(self.module_offloaders),
            self.group_size,
            self.num_in_group,
        )
        return modules

    def post_init(self) -> None:
        self.total_offloaded_bytes = 0
        for module_offloader in self.module_offloaders:
            module_offloader.post_init()
            self.total_offloaded_bytes += module_offloader.offloaded_bytes
        logger.info(
            "[AscendSyncOffloader] initialized wrapped=%d, saved=%.4f GB, "
            "selected_layers=%s",
            len(self.module_offloaders),
            self.total_offloaded_bytes / float(2**30),
            self.selected_layer_indices,
        )
