import contextlib
import math
from typing import Dict, Iterable, List, Sequence

import torch
import torch.nn as nn


class LoRALinear(nn.Module):
    """
    Lightweight LoRA wrapper for nn.Linear.
    """

    def __init__(
        self,
        base_linear: nn.Linear,
        rank: int = 16,
        alpha: float = 16.0,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        if not isinstance(base_linear, nn.Linear):
            raise TypeError("LoRALinear expects an nn.Linear.")
        if rank <= 0:
            raise ValueError("LoRA rank must be > 0.")

        self.base = base_linear
        self.rank = int(rank)
        self.alpha = float(alpha)
        self.scaling = self.alpha / float(self.rank)
        self.enabled = True
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

        self.lora_A = nn.Linear(self.base.in_features, self.rank, bias=False)
        self.lora_B = nn.Linear(self.rank, self.base.out_features, bias=False)
        self.reset_parameters()
        self.base.requires_grad_(False)

    def reset_parameters(self) -> None:
        nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.base(x)
        if not self.enabled:
            return out
        delta = self.lora_B(self.lora_A(self.dropout(x)))
        return out + self.scaling * delta


def _match_target(name: str, target_modules: Sequence[str]) -> bool:
    leaf = name.split(".")[-1]
    for target in target_modules:
        if not target:
            continue
        if leaf == target or name.endswith(target):
            return True
    return False


def _find_parent_module(root: nn.Module, module_name: str):
    parts = module_name.split(".")
    parent = root
    for part in parts[:-1]:
        parent = getattr(parent, part)
    return parent, parts[-1]


def inject_lora(
    model: nn.Module,
    target_modules: Sequence[str],
    rank: int,
    alpha: float,
    dropout: float,
) -> List[str]:
    named_modules = list(model.named_modules())
    injected = []
    for module_name, module in named_modules:
        if not _match_target(module_name, target_modules):
            continue
        if isinstance(module, LoRALinear):
            continue
        if not isinstance(module, nn.Linear):
            continue
        parent, child_name = _find_parent_module(model, module_name)
        setattr(parent, child_name, LoRALinear(module, rank=rank, alpha=alpha, dropout=dropout))
        injected.append(module_name)
    return injected


def has_lora(model: nn.Module) -> bool:
    return any(isinstance(m, LoRALinear) for m in model.modules())


def lora_parameters(model: nn.Module) -> Iterable[nn.Parameter]:
    for module in model.modules():
        if isinstance(module, LoRALinear):
            yield from module.lora_A.parameters()
            yield from module.lora_B.parameters()


def freeze_non_lora_parameters(model: nn.Module) -> None:
    model.requires_grad_(False)
    for p in lora_parameters(model):
        p.requires_grad_(True)


def lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]:
    state = {}
    for name, module in model.named_modules():
        if not isinstance(module, LoRALinear):
            continue
        state[f"{name}.lora_A.weight"] = module.lora_A.weight.detach().cpu()
        state[f"{name}.lora_B.weight"] = module.lora_B.weight.detach().cpu()
    return state


def load_lora_state_dict(model: nn.Module, state_dict: Dict[str, torch.Tensor], strict: bool = True) -> None:
    missing = []
    for name, module in model.named_modules():
        if not isinstance(module, LoRALinear):
            continue
        key_a = f"{name}.lora_A.weight"
        key_b = f"{name}.lora_B.weight"
        if key_a not in state_dict or key_b not in state_dict:
            missing.append(name)
            continue
        module.lora_A.weight.data.copy_(state_dict[key_a].to(module.lora_A.weight.device))
        module.lora_B.weight.data.copy_(state_dict[key_b].to(module.lora_B.weight.device))
    if strict and missing:
        raise KeyError(f"Missing LoRA weights for modules: {missing[:8]}")


def set_lora_enabled(model: nn.Module, enabled: bool) -> None:
    for module in model.modules():
        if isinstance(module, LoRALinear):
            module.enabled = bool(enabled)


@contextlib.contextmanager
def lora_disabled(model: nn.Module):
    set_lora_enabled(model, enabled=False)
    try:
        yield
    finally:
        set_lora_enabled(model, enabled=True)