import torch
import torch.nn as nn
import torch.nn.functional as F
from ..utils.utils import match_target_modules, get_parent_module
from ..utils.introspect import get_module_input_output_dims, has_valid_weight

ADAPTER_BLACKLIST = {"q_proj", "k_proj", "v_proj", "o_proj", "c_attn", "query", "key", "value", "out"}

class Adapter(nn.Module):
    def __init__(self, hidden_dim, bottleneck):
        super().__init__()
        self.down = nn.Linear(hidden_dim, bottleneck)
        self.up = nn.Linear(bottleneck, hidden_dim)
        nn.init.zeros_(self.up.weight)
        nn.init.zeros_(self.up.bias)

    def forward(self, x):
        return x + self.up(F.relu(self.down(x)))

class AdapterWrapper(nn.Module):
    def __init__(self, orig_linear, adapter):
        super().__init__()
        self.orig = orig_linear
        self.adapter = adapter
        for p in self.orig.parameters():
            p.requires_grad = False

    def forward(self, x):
        return self.adapter(self.orig(x))

class Adapters:
    def __init__(self, model, bottleneck, target_modules=None, match_mode='contains'):
        self.model = model
        self.adapters = []
        if target_modules is None:
            target_modules = ['mlp', 'fc1', 'intermediate', 'dense_h_to_4h', 'linear', 'fc']
        exclude = ['norm', 'dropout', 'embed']
        for p in self.model.parameters():
            p.requires_grad = False
        candidates = [
            (name, module)
            for name, module in self.model.named_modules()
            if has_valid_weight(module)
        ]
        for name, module in candidates:
            if any(pattern in name.lower() for pattern in exclude):
                continue
            if self._should_skip(name):
                continue
            if not match_target_modules(name, target_modules, mode=match_mode):
                continue
            try:
                _, out_dim = get_module_input_output_dims(module)
                if out_dim is None:
                    continue
                parent, attr = get_parent_module(self.model, name)
                adapter = Adapter(out_dim, bottleneck)
                wrapper = AdapterWrapper(module, adapter)
                setattr(parent, attr, wrapper)
                self.adapters.append(adapter)
            except Exception:
                continue

    def _should_skip(self, module_name):
        last_component = module_name.split(".")[-1].lower()
        return any(blacklisted in last_component for blacklisted in ADAPTER_BLACKLIST)

    def parameters(self):
        for adapter in self.adapters:
            for param in adapter.parameters():
                yield param

def inject_adapter(model, bottleneck=64, target_modules=None, match_mode='contains'):
    if target_modules is None:
        target_modules = ['mlp', 'fc1', 'intermediate', 'dense_h_to_4h', 'linear', 'fc']
    peft = Adapters(model, bottleneck, target_modules, match_mode)
    return model, peft