"""Gated LoRA adapters of Vela 3.0.

The adapters (vela3_config.json "adapters") sit on the base's Linear layers; each adapter's layers, rank and alpha come from its
*_lora.safetensors. At most one adapter is on in any forward; with no adapter active, a wrapped layer returns base(x) itself, so every
other pass is the untouched base computation.

The adapter math is the training code's GatedLoRA in eval mode (dropout off):
    y = base(x) + ((x.float() @ A.T) @ B.T * alpha / r).to(y.dtype)
A row mask applies the adapter to some rows of a batch only: torch.where keeps base(x) bit for bit on the other rows."""
from __future__ import annotations

import contextlib
import json
import threading

import torch
from torch import nn


class _Gate(threading.local):  # per thread: concurrent requests never see each other's adapter
    name = None   # active adapter name, or None (all adapters off)
    rows = None   # None = every row of the batch; else a bool tensor [B]


GATE = _Gate()


@contextlib.contextmanager
def active(name, rows=None):
    """Activate one adapter for the forwards inside the block (rows: optional bool mask over the batch)."""
    prev = (GATE.name, GATE.rows)
    GATE.name, GATE.rows = name, rows
    try:
        yield
    finally:
        GATE.name, GATE.rows = prev


class MultiLoRA(nn.Module):
    def __init__(self, base: nn.Linear):
        super().__init__()
        self.base = base
        self.ad = {}  # name -> (A [r, in] fp32, B [out, r] fp32, scale)

    def add(self, name, A, B, scale):
        dev = self.base.weight.device
        self.ad[name] = (A.to(dev, torch.float32), B.to(dev, torch.float32), float(scale))

    def forward(self, x):
        y = self.base(x)
        name = GATE.name
        if name is None or name not in self.ad:
            return y
        A, B, s = self.ad[name]
        d = ((x.float() @ A.T) @ B.T * s).to(y.dtype)
        rows = GATE.rows
        if rows is None:
            return y + d
        if rows.shape[0] != y.shape[0]:
            raise RuntimeError("row gate does not match the batch")
        return torch.where(rows.view(-1, *([1] * (y.dim() - 1))), y + d, y)


def attach(lm, name, path):
    """Attach the adapter in `path` (safetensors: '<module path>.A' / '.B', metadata r / alpha / first / targets) to the text model."""
    from safetensors import safe_open
    with safe_open(str(path), framework="pt") as f:
        meta = f.metadata() or {}
        sd = {k: f.get_tensor(k) for k in f.keys()}
    r, alpha = int(meta["r"]), int(meta["alpha"])
    mods = sorted({k.rsplit(".", 1)[0] for k in sd})
    for mp in mods:
        parent_path, attr = mp.rsplit(".", 1)
        parent = lm.get_submodule(parent_path)
        cur = getattr(parent, attr)
        if isinstance(cur, nn.Linear):
            cur = MultiLoRA(cur)
            setattr(parent, attr, cur)
        elif not isinstance(cur, MultiLoRA):
            raise ValueError(f"{mp} is not a Linear layer")
        A, B = sd[mp + ".A"], sd[mp + ".B"]
        if A.shape != (r, cur.base.in_features) or B.shape != (cur.base.out_features, r):
            raise ValueError(f"adapter shape mismatch at {mp}")
        cur.add(name, A, B, alpha / r)
    if len(sd) != 2 * len(mods):
        raise ValueError(f"adapter {name}: unpaired A / B tensors")
    return dict(modules=len(mods), r=r, alpha=alpha, first=int(meta.get("first", -1)),
                targets=json.loads(meta.get("targets", "[]")))