from contextlib import contextmanager
import importlib
import inspect
import sys

import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.models.utils import extract_layer_index
from vllm.model_executor.models import *


@contextmanager
def patch(module_path: str, new_cls_dict: dict[str, type]):
  """Helper function for patching into model implementations

  Args:
      module_path (str): the module path to patch
      new_cls_dict (dict[str, type]): all classes to patch, and the new classes
        to patch in.
  """
  module = sys.modules[module_path]
  old_values = {a: getattr(module, a) for a in new_cls_dict.keys()}

  try:
    for a, v in new_cls_dict.items():
      setattr(module, a, v)
    yield
  finally:
    for a, v in old_values.items():
      setattr(module, a, v)


def build_ablated_module_class(
    module_path: str,
    orig_cls_name: str,
) -> type:
  """Dynamically subclasses attention/mlp classes

  Args:
      module_path (str): the module path where the original class is used,
        e.g. 'vllm.model_executor.models.qwen3_5'
      orig_cls_name (str): the name of the original class to subclass,
        e.g. 'GatedDeltaNetAttention'

  Raises:
      ValueError: Thrown if an unknown ablation is passed

  Returns:
      type: the patched subclass
  """
  orig_cls = getattr(sys.modules[module_path], orig_cls_name)
  orig_forward = orig_cls.forward

  # Introspect the original forward signature once at build time.
  sig = inspect.signature(orig_forward)
  param_names = [n for n in sig.parameters if n != 'self']

  # Find the hidden-states parameter: prefer 'hidden_states', then 'x', then first param.
  hs_candidates = ['hidden_states', 'x']
  hs_param = next((p for p in hs_candidates if p in param_names), None)
  if hs_param is None and param_names:
    hs_param = param_names[0]

  # Build a signature without 'self' for call-site binding.
  non_self_params = [p for n, p in sig.parameters.items() if n != 'self']
  call_sig = sig.replace(parameters=non_self_params)

  # Detect output-buffer convention (writes into 'output', returns None).
  has_output = 'output' in param_names

  def forward(self, *args, **kwargs):
    ablation_impl = getattr(self, 'ablation_implementation', None)

    if ablation_impl is None:
      return orig_forward(self, *args, **kwargs)

    bound = call_sig.bind(*args, **kwargs)
    bound.apply_defaults()
    hs = bound.arguments[hs_param]
    out = bound.arguments.get('output') if has_output else None

    match ablation_impl:
      case 'zero':
        if out is not None:
          out.zero_()
          return
        return torch.zeros_like(hs)
      case 'identity':
        if out is not None:
          out.copy_(hs)
          return
        return hs
      case _:
        raise ValueError(f'Unknown ablation implementation: {ablation_impl}')

  return type(
      f'Ablated{orig_cls_name}',
      (orig_cls,),
      {'forward': forward},
  )


def build_ablated_decoder_layer(
    module_path: str,
    orig_cls_name: str,
    ablated_attn_class_names: set[str],
    ablated_mlp_class_names: set[str],
) -> type:
  """Dynamically subclasses the decoder layer class

  Args:
      module_path (str): the module path where the original class is used,
        e.g. 'vllm.model_executor.models.qwen3_5'
      orig_cls_name (str): the name of the original class to subclass,
        e.g. 'Qwen3_5DecoderLayer'
      ablated_attn_class_names (set[str]): the set of attention class names
        that are ablated, e.g.
      ablated_mlp_class_names (set[str]): the set of MLP class names
        that are ablated, e.g.

  Returns:
      type: the patched decoder layer subclass
  """
  orig_cls = getattr(sys.modules[module_path], orig_cls_name)

  def __init__(
      self: nn.Module,
      vllm_config: VllmConfig,
      *args,
      prefix: str = '',
      **kwargs,
  ):
    orig_cls.__init__(
        self,
        vllm_config=vllm_config,
        *args,
        prefix=prefix,
        **kwargs,
    )
    layer_index = str(extract_layer_index(prefix))

    config = vllm_config.model_config.hf_config
    ablated_attention_layers = getattr(config, 'ablated_attention_layers', {})
    ablated_mlp_layers = getattr(config, 'ablated_mlp_layers', {})

    ablated_attn_impl = ablated_attention_layers.get(layer_index, None)
    ablated_mlp_impl = ablated_mlp_layers.get(layer_index, None)

    # Set the ablation implementation on submodules of this decoder layer.
    for module in self.children():
      class_name = module.__class__.__name__
      if class_name in ablated_attn_class_names:
        module.ablation_implementation = ablated_attn_impl
      elif class_name in ablated_mlp_class_names:
        module.ablation_implementation = ablated_mlp_impl

  return type(
      f'Ablated{orig_cls_name}',
      (orig_cls,),
      {'__init__': __init__},
  )


def build_ablated_model_class(
    causal_lm_class: str,
    decoder_class: str,
    attention_classes: list[str],
    mlp_classes: list[str],
    module_path: str,
    additional_subclasses: list[str] | None = None,
) -> type:
  """Dynamically subclass a Causal LM class to support ablations

  Args:
      causal_lm_class (str): the name of the original causal LM class to subclass,
        e.g. 'Qwen3_5ForConditionalGeneration'
      decoder_class (str): the name of the decoder class to subclass,
        e.g. 'Qwen3_5DecoderLayer'
      attention_classes (list[str]): the list of attention class names to ablate,
        e.g. ['Qwen3_5Attention']
      mlp_classes (list[str]): the list of MLP class names to ablate,
        e.g. ['Qwen3_5MLP']
      module_path (str): the module path where the original classes are used,
        e.g. 'vllm.model_executor.models.qwen3_5'
      additional_subclasses (list[str] | None, optional): the list of additional
        subclasses to include. Defaults to None.

  Returns:
      type: the patched causal LM subclass, to be registered in the ModelRegistry
  """
  if additional_subclasses is None:
    additional_subclasses = []

  importlib.import_module(module_path)

  # build ablated attention and mlp classes
  ablated_attention_classes = {
      cls_name: build_ablated_module_class(module_path, cls_name)
      for cls_name in attention_classes
  }
  ablated_mlp_classes = {
      cls_name: build_ablated_module_class(module_path, cls_name)
      for cls_name in mlp_classes
  }

  # build ablated decoder layer
  ablated_attn_class_names = set([f'Ablated{n}' for n in attention_classes])
  ablated_mlp_class_names = set([f'Ablated{n}' for n in mlp_classes])

  ablated_decoder_layer = build_ablated_decoder_layer(
      module_path=module_path,
      orig_cls_name=decoder_class,
      ablated_attn_class_names=ablated_attn_class_names,
      ablated_mlp_class_names=ablated_mlp_class_names,
  )

  # get original causal LM class
  orig_causal_lm_cls = getattr(sys.modules[module_path], causal_lm_class)

  subclasses = []
  if additional_subclasses:
    for subclass_path in additional_subclasses:
      module_name, class_name = subclass_path.rsplit('.', 1)
      module = importlib.import_module(module_name)
      subclasses.append(getattr(module, class_name))

  # init function where decoder layers & attention/mlp classes are patched in
  def __init__(self, *, vllm_config: VllmConfig, prefix: str):
    with patch(
        module_path,
        {
            **ablated_attention_classes,
            **ablated_mlp_classes,
            decoder_class: ablated_decoder_layer,
        },
    ):
      orig_causal_lm_cls.__init__(self, vllm_config=vllm_config, prefix=prefix)

  return type(
      f'Ablated{causal_lm_class}',
      (orig_causal_lm_cls, *subclasses),
      {'__init__': __init__},
  )
