import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Tuple

from ..utils.introspect import get_module_input_output_dims, get_device, get_dtype
from ..utils.utils import match_target_modules, get_parent_module


class LoRALayer(nn.Module):
    def __init__(self, base: nn.Module, r: int = 8, alpha: int = 16, dropout: float = 0.0):
        super().__init__()
        self.base = base
        self.r = r
        self.scaling = alpha / max(1, r)
        self.dropout = dropout
        self.merged = False

        in_dim, out_dim = get_module_input_output_dims(base)
        if in_dim is None or out_dim is None:
            raise ValueError(f"Cannot determine dims for {type(base).__name__}")

        device, dtype = get_device(base), get_dtype(base)

        if r > 0:
            init_scale = 0.01
            self.A = nn.Parameter(torch.randn(r, in_dim, device=device, dtype=dtype) * init_scale)
            self.B = nn.Parameter(torch.randn(out_dim, r, device=device, dtype=dtype) * init_scale)
        else:
            self.register_parameter("A", None)
            self.register_parameter("B", None)

        for p in self.base.parameters():
            p.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.merged:
            return self.base(x)
        base_out = self.base(x)
        if self.r <= 0 or self.A is None:
            return base_out
        x = F.dropout(x, p=self.dropout, training=self.training)
        lora_out = x @ self.A.T @ self.B.T
        return base_out + lora_out * self.scaling


def inject_lora(
    model: nn.Module,
    r: int = 8,
    alpha: int = 16,
    target_modules: Optional[List[str]] = None,
    exclude_modules: Optional[List[str]] = None,
    match_mode: str = "contains",
    dropout: float = 0.0,
) -> Tuple[nn.Module, List[LoRALayer]]:
    if target_modules is None:
        target_modules = ["linear", "fc"]
    if exclude_modules is None:
        exclude_modules = ["norm", "embed", "dropout"]

    for p in model.parameters():
        p.requires_grad = False

    lora_layers = []

    from ..utils.introspect import get_trainable_modules
    candidates = get_trainable_modules(model, exclude_patterns=exclude_modules)

    for name, module in candidates:
        module_key = f"{name}.{module.__class__.__name__}"
        if not match_target_modules(module_key, target_modules, mode=match_mode):
            continue
        try:
            parent, attr = get_parent_module(model, name)
            lora_layer = LoRALayer(module, r=r, alpha=alpha, dropout=dropout)
            setattr(parent, attr, lora_layer)
            lora_layers.append(lora_layer)
        except Exception:
            continue

    return model, lora_layers


def get_lora_state_dict(model: nn.Module) -> dict:
    sd = {}
    for name, module in model.named_modules():
        if isinstance(module, LoRALayer) and module.r > 0:
            sd[f"{name}.A"] = module.A
            sd[f"{name}.B"] = module.B
    return sd


def load_lora_state_dict(model: nn.Module, state: dict) -> nn.Module:
    for name, module in model.named_modules():
        if isinstance(module, LoRALayer):
            if f"{name}.A" in state:
                module.A.data = state[f"{name}.A"]
            if f"{name}.B" in state:
                module.B.data = state[f"{name}.B"]
    return model


def merge_lora_weights(model: nn.Module) -> nn.Module:
    with torch.no_grad():
        for module in model.modules():
            if isinstance(module, LoRALayer) and module.r > 0 and isinstance(module.base, nn.Linear):
                delta = module.B @ module.A
                module.base.weight += delta * module.scaling
                module.merged = True
    print("[merge_lora_weights] merged")
    return model


def unmerge_lora_weights(model: nn.Module) -> nn.Module:
    with torch.no_grad():
        for module in model.modules():
            if isinstance(module, LoRALayer) and module.r > 0 and isinstance(module.base, nn.Linear):
                delta = module.B @ module.A
                module.base.weight -= delta * module.scaling
                module.merged = False
    print("[unmerge_lora_weights] unmerged")
    return model