"""Attach CAFE-25lang's language-conditional adapters to a restored NeMo model.

The .nemo export carries the trunk only, which is the English path. Every other language needs the
adapters from `adapters.safetensors` re-attached, which this helper does with residual forward
hooks so the checkpoint itself is never modified.

Used by eval_mls8.py (`--adapters ...`) and by the official Open ASR Leaderboard runners through
run_leaderboard.py.
"""
from __future__ import annotations

import torch
import torch.nn as nn
from safetensors.torch import load_file

ENCODER_LAYERS = (3, 7, 11, 15, 19, 23)


class Bottleneck(nn.Module):
    """x -> x + up(gelu(down(x)))"""

    def __init__(self, dim: int, rank: int):
        super().__init__()
        self.down = nn.Linear(dim, rank)
        self.act = nn.GELU()
        self.up = nn.Linear(rank, dim)

    def forward(self, x):
        return self.up(self.act(self.down(x)))


def attach_adapters(model, path: str, device, dtype=torch.bfloat16, verbose: bool = True):
    """Inject the adapters; returns the hook handles so they can be removed again."""
    sd = load_file(path)
    idxs = sorted({int(k.split(".")[1]) for k in sd})
    rank, dim = sd["adapter.0.down.weight"].shape
    if len(idxs) != len(ENCODER_LAYERS):
        raise SystemExit(f"{len(idxs)} adapters for {len(ENCODER_LAYERS)} encoder layers")
    handles, mods = [], []
    for i, layer in zip(idxs, ENCODER_LAYERS):
        a = Bottleneck(dim, rank)
        a.load_state_dict({"down.weight": sd[f"adapter.{i}.down.weight"],
                           "down.bias": sd[f"adapter.{i}.down.bias"],
                           "up.weight": sd[f"adapter.{i}.up.weight"],
                           "up.bias": sd[f"adapter.{i}.up.bias"]})
        a = a.to(device=device, dtype=dtype).eval()
        mods.append(a)

        def mk(a):
            def hook(module, inp, out):
                x = out[0] if isinstance(out, tuple) else out
                if x.dim() != 3 or x.shape[-1] != a.down.in_features:
                    return out          # shape guard: leave the graph untouched
                y = x + a(x)
                return (y,) + tuple(out[1:]) if isinstance(out, tuple) else y
            return hook

        handles.append(model.encoder.layers[layer].register_forward_hook(mk(a)))
    model._adapters, model._adapter_handles = mods, handles
    if verbose:
        print(f"[adapters] attached at encoder layers {list(ENCODER_LAYERS)} "
              f"(dim={dim}, rank={rank})")
    return handles