"""Resident dynamic LoRA branches for the quantized H3 denoiser and refiner."""

from collections.abc import Iterator

import torch
import torch.nn.functional as functional
from safetensors import safe_open
from torch import nn


class LoraBranch(nn.Module):
    def __init__(self, down: torch.Tensor, up: torch.Tensor, alpha: float):
        super().__init__()
        if down.ndim != 2 or up.ndim != 2 or down.shape[0] != up.shape[1]:
            raise ValueError(f"Invalid LoRA shapes: down={tuple(down.shape)}, up={tuple(up.shape)}")
        self.scale = float(alpha) / down.shape[0]
        self.register_buffer("down", down.contiguous(), persistent=False)
        self.register_buffer("up", up.contiguous(), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.to(self.down.dtype)
        return functional.linear(functional.linear(x, self.down), self.up) * self.scale


class DynamicLoraMixin:
    """Small mixin for linears that can host resident request-selectable LoRAs."""

    def _init_dynamic_lora(self) -> None:
        self.lora_branches = nn.ModuleDict()
        self.active_lora: str | None = None
        self.lora_strength = 0.0

    def add_lora(self, name: str, down: torch.Tensor, up: torch.Tensor, alpha: float) -> None:
        if name in self.lora_branches:
            raise ValueError(f"LoRA {name!r} is already attached")
        if down.shape[1] != self.in_features or up.shape[0] != self.out_features:
            raise ValueError(
                f"LoRA {name!r} dimensions {tuple(down.shape)}, {tuple(up.shape)} do not match "
                f"linear [{self.out_features}, {self.in_features}]"
            )
        self.lora_branches[name] = LoraBranch(down, up, alpha)

    def set_lora(self, name: str | None, strength: float = 1.0) -> None:
        if name is not None and name not in self.lora_branches:
            raise ValueError(f"LoRA {name!r} is not attached")
        self.active_lora = name
        self.lora_strength = float(strength) if name is not None else 0.0

    def _apply_lora(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
        if self.active_lora is None or self.lora_strength == 0.0:
            return base
        delta = self.lora_branches[self.active_lora](x)
        return base + delta.to(base.dtype) * self.lora_strength


def iter_lora_targets(model: nn.Module, refiner: nn.Module) -> Iterator[tuple[str, DynamicLoraMixin]]:
    for index, block in enumerate(model.backbone.blocks):
        yield f"blocks.{index}.attn.qkv_proj", block.attention.qkv_proj
        yield f"blocks.{index}.attn.out_proj", block.attention.out_proj
        yield f"blocks.{index}.mlp.fc1", block.mlp.fc1
        yield f"blocks.{index}.mlp.fc2", block.mlp.fc2
    for index, block in enumerate(refiner.blocks):
        yield f"token_refiner.blocks.{index}.attn.qkv_proj", block.qkv
        yield f"token_refiner.blocks.{index}.attn.out_proj", block.out
        yield f"token_refiner.blocks.{index}.mlp.fc1", block.fc1
        yield f"token_refiner.blocks.{index}.mlp.fc2", block.fc2


def load_lora_adapter(model: nn.Module, refiner: nn.Module, name: str, path: str, device: str) -> int:
    targets = list(iter_lora_targets(model, refiner))
    expected = {
        f"diffusion_model.{target}.{suffix}"
        for target, _module in targets
        for suffix in ("alpha", "lora_A.weight", "lora_B.weight")
    }
    with safe_open(path, framework="pt", device=device) as checkpoint:
        actual = set(checkpoint.keys())
        if actual != expected:
            missing = sorted(expected - actual)[:8]
            unexpected = sorted(actual - expected)[:8]
            raise ValueError(f"LoRA key mismatch: missing={missing}, unexpected={unexpected}")
        for target, module in targets:
            prefix = f"diffusion_model.{target}"
            module.add_lora(
                name,
                checkpoint.get_tensor(f"{prefix}.lora_A.weight"),
                checkpoint.get_tensor(f"{prefix}.lora_B.weight"),
                checkpoint.get_tensor(f"{prefix}.alpha").item(),
            )
    return len(targets)


def set_active_lora(model: nn.Module, refiner: nn.Module, name: str | None, strength: float = 1.0) -> None:
    for _target, module in iter_lora_targets(model, refiner):
        module.set_lora(name, strength)