"""v0.71.16 — Knowledge edit depth.

Closes:
  * #251 — edit kernels gain GPT-2 ``transformer.h`` / ``mlp.c_proj`` support
           (transpose-aware rank-1 update for the Conv1D weight layout).
  * #252 — EditGovernor edit-count increment is now atomic (baseline-delta
           merge under the cross-process lock) so concurrent ``soup edit set``
           runs cannot lose an increment.
  * #250 — covariance-preconditioned ROME via ``--cov-corpus`` (estimate the
           key covariance C from a stats corpus and use C^{-1} k* instead of
           k*; falls back to C=I when no corpus).
  * #147 — Mixtral joins the LongLoRA architecture allowlist (dedicated
           ``is_mixtral_model`` helper + ``MixtralAttention`` forward override).

Kernel maths are exercised with real torch (the [dev] extra) on tiny CPU
fakes. Full real-model apply paths are covered by the release step-6 smoke
(tiny-gpt2 + SmolLM2-135M).
"""

from __future__ import annotations

import json
import os
import sys
from types import SimpleNamespace

import pytest

torch = pytest.importorskip("torch")
nn = torch.nn  # assignment (not an import) so it works after importorskip


# ===========================================================================
# Shared fakes — a faithful tiny GPT-2 / Llama LM (forward + loss + tokenizer)
# ===========================================================================


class _Conv1D(nn.Module):
    """Faithful mini ``transformers.pytorch_utils.Conv1D``.

    Weight is ``[nx, nf]`` (transposed relative to nn.Linear's ``[out, in]``),
    ``nf`` is the output feature count, and ``forward`` is ``x @ weight + bias``.
    """

    def __init__(self, nx: int, nf: int):
        super().__init__()
        self.nf = nf
        self.weight = nn.Parameter(torch.randn(nx, nf) * 0.05)
        self.bias = nn.Parameter(torch.zeros(nf))

    def forward(self, x):  # noqa: D401
        return x @ self.weight + self.bias


class _FakeEnc(dict):
    """BatchEncoding-like: dict for ``**`` unpacking + a no-op ``.to``."""

    def to(self, _device):
        return self


class _FakeTok:
    vocab = 24

    def _ids(self, text: str) -> list[int]:
        out = [(sum(ord(c) for c in t) % (self.vocab - 4)) + 1 for t in text.split()]
        return out or [1]

    def __call__(
        self,
        text,
        return_tensors=None,
        truncation=False,
        max_length=None,
        add_special_tokens=True,
    ):
        ids = self._ids(text)
        if add_special_tokens:
            ids = [0] + ids
        if max_length is not None:
            ids = ids[:max_length]
        if not ids:
            ids = [1]
        if return_tensors == "pt":
            return _FakeEnc(input_ids=torch.tensor([ids], dtype=torch.long))
        return {"input_ids": ids}


def _lm_forward(hidden, lm_head, input_ids, labels):
    logits = lm_head(hidden)
    out = SimpleNamespace(logits=logits)
    if labels is not None:
        shift_logits = logits[:, :-1, :].reshape(-1, logits.shape[-1])
        shift_labels = labels[:, 1:].reshape(-1)
        out.loss = nn.functional.cross_entropy(
            shift_logits, shift_labels, ignore_index=-100
        )
    return out


class _LlamaMLP(nn.Module):
    def __init__(self, hidden, inter):
        super().__init__()
        self.gate_proj = nn.Linear(hidden, inter, bias=False)
        self.down_proj = nn.Linear(inter, hidden, bias=False)

    def forward(self, x):
        return self.down_proj(torch.relu(self.gate_proj(x)))


class _LlamaBlock(nn.Module):
    def __init__(self, hidden, inter):
        super().__init__()
        self.mlp = _LlamaMLP(hidden, inter)

    def forward(self, x):
        return x + self.mlp(x)


class _LlamaInner(nn.Module):
    def __init__(self, vocab, hidden, inter, layers):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab, hidden)
        self.layers = nn.ModuleList(
            [_LlamaBlock(hidden, inter) for _ in range(layers)]
        )

    def forward(self, input_ids):
        x = self.embed_tokens(input_ids)
        for blk in self.layers:
            x = blk(x)
        return x


class _FakeLlamaLM(nn.Module):
    def __init__(self, vocab=24, hidden=8, inter=16, layers=3):
        super().__init__()
        self.model = _LlamaInner(vocab, hidden, inter, layers)
        self.lm_head = nn.Linear(hidden, vocab, bias=False)

    def forward(self, input_ids=None, labels=None, **_kw):
        return _lm_forward(self.model(input_ids), self.lm_head, input_ids, labels)


class _GPT2MLP(nn.Module):
    def __init__(self, hidden, inter):
        super().__init__()
        self.c_fc = _Conv1D(hidden, inter)
        self.c_proj = _Conv1D(inter, hidden)

    def forward(self, x):
        return self.c_proj(torch.relu(self.c_fc(x)))


class _GPT2Block(nn.Module):
    def __init__(self, hidden, inter):
        super().__init__()
        self.mlp = _GPT2MLP(hidden, inter)

    def forward(self, x):
        return x + self.mlp(x)


class _GPT2Inner(nn.Module):
    def __init__(self, vocab, hidden, inter, layers):
        super().__init__()
        self.wte = nn.Embedding(vocab, hidden)
        self.h = nn.ModuleList([_GPT2Block(hidden, inter) for _ in range(layers)])

    def forward(self, input_ids):
        x = self.wte(input_ids)
        for blk in self.h:
            x = blk(x)
        return x


class _FakeGPT2LM(nn.Module):
    def __init__(self, vocab=24, hidden=8, inter=16, layers=3):
        super().__init__()
        self.transformer = _GPT2Inner(vocab, hidden, inter, layers)
        self.lm_head = nn.Linear(hidden, vocab, bias=False)

    def forward(self, input_ids=None, labels=None, **_kw):
        return _lm_forward(
            self.transformer(input_ids), self.lm_head, input_ids, labels
        )


def _peft_wrap(base):
    """Minimal PEFT-style wrapper exposing ``get_base_model``."""

    class _Peft:
        def __init__(self, b):
            self._b = b

        def get_base_model(self):
            return self._b

    return _Peft(base)


class _PeftCallable:
    """Callable PEFT-style wrapper: delegates forward + train/eval to base.

    Unlike ``_peft_wrap`` (locate-only), this supports the full kernel path —
    ``model(**inputs)`` / ``model.training`` / ``eval()`` / ``train()`` — so a
    GPT-2 base can be edited end-to-end while wrapped (L1).
    """

    def __init__(self, base):
        self._b = base

    def get_base_model(self):
        return self._b

    def __call__(self, *a, **k):
        return self._b(*a, **k)

    @property
    def training(self):
        return self._b.training

    def eval(self):
        self._b.eval()
        return self

    def train(self, mode=True):
        self._b.train(mode)
        return self


_SUBJECT = "Paris is the capital of"
_TARGET = "Lyon"
_CORPUS = [
    "the quick brown fox jumps over the lazy dog",
    "lorem ipsum dolor sit amet consectetur",
    "machine learning models edit facts surgically",
]


# ===========================================================================
# #251 — GPT-2 transformer.h / mlp.c_proj support
# ===========================================================================


class TestLocateDecoderLayersGpt2:
    def test_locates_transformer_h(self):
        from soup_cli.utils.edit_kernels import _locate_decoder_layers

        model = _FakeGPT2LM(layers=4)
        layers = _locate_decoder_layers(model)
        assert len(layers) == 4

    def test_peft_wrapped_gpt2(self):
        from soup_cli.utils.edit_kernels import _locate_decoder_layers

        model = _peft_wrap(_FakeGPT2LM(layers=3))
        layers = _locate_decoder_layers(model)
        assert len(layers) == 3

    def test_peft_wrapped_llama_still_works(self):
        from soup_cli.utils.edit_kernels import _locate_decoder_layers

        model = _peft_wrap(_FakeLlamaLM(layers=2))
        layers = _locate_decoder_layers(model)
        assert len(layers) == 2

    def test_unknown_arch_raises(self):
        from soup_cli.utils.edit_kernels import _locate_decoder_layers

        with pytest.raises(ValueError, match="decoder layers"):
            _locate_decoder_layers(nn.Linear(2, 2))

    def test_get_base_model_raises_swallowed(self):
        """A PEFT wrapper whose get_base_model() blows up must fall through to
        the clear ValueError (DEBUG-logged, not masked) — review L2."""
        from soup_cli.utils.edit_kernels import _locate_decoder_layers

        class _BadPeft:
            def get_base_model(self):
                raise RuntimeError("boom")

        with pytest.raises(ValueError, match="decoder layers"):
            _locate_decoder_layers(_BadPeft())


class TestDownProjGpt2:
    def test_returns_c_proj(self):
        from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers

        model = _FakeGPT2LM()
        down = _down_proj(_locate_decoder_layers(model), 1)
        assert hasattr(down, "weight")
        # Conv1D weight is [in, out] = [inter, hidden] = [16, 8].
        assert tuple(down.weight.shape) == (16, 8)
        assert down.nf == 8

    def test_llama_down_proj_unchanged(self):
        from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers

        model = _FakeLlamaLM()
        down = _down_proj(_locate_decoder_layers(model), 0)
        # nn.Linear weight is [out, in] = [hidden, inter] = [8, 16].
        assert tuple(down.weight.shape) == (8, 16)

    def test_out_of_range(self):
        from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers

        layers = _locate_decoder_layers(_FakeGPT2LM())
        with pytest.raises(ValueError, match="out of range"):
            _down_proj(layers, 99)


class TestProjHelpers:
    def test_is_transposed_proj_conv1d(self):
        from soup_cli.utils.edit_kernels import _is_transposed_proj

        assert _is_transposed_proj(_Conv1D(16, 8)) is True

    def test_is_transposed_proj_linear(self):
        from soup_cli.utils.edit_kernels import _is_transposed_proj

        assert _is_transposed_proj(nn.Linear(16, 8, bias=False)) is False

    def test_proj_out_dim_conv1d_uses_nf(self):
        from soup_cli.utils.edit_kernels import _proj_out_dim

        # Conv1D nf = output (hidden) dim, NOT weight.shape[0] (= in dim).
        assert _proj_out_dim(_Conv1D(16, 8)) == 8

    def test_proj_out_dim_linear_uses_shape0(self):
        from soup_cli.utils.edit_kernels import _proj_out_dim

        assert _proj_out_dim(nn.Linear(16, 8, bias=False)) == 8

    def test_is_transposed_proj_rejects_bool_nf(self):
        """``nf=True`` (bool, a subclass of int) must NOT be treated as Conv1D
        — review L3."""
        from soup_cli.utils.edit_kernels import _is_transposed_proj

        assert _is_transposed_proj(SimpleNamespace(nf=True)) is False


class TestRank1UpdateTransposed:
    def test_conv1d_post_condition(self):
        """Conv1D: ``key @ W`` must gain exactly ``delta`` after the update."""
        from soup_cli.utils.edit_kernels import _rank1_update

        conv = _Conv1D(16, 8)
        key = torch.ones(16)
        delta = torch.full((8,), 0.5)
        before = key @ conv.weight  # [out] = [8]
        norm = _rank1_update(conv, key, delta)
        after = key @ conv.weight
        assert norm > 0
        assert torch.allclose(after - before, delta, atol=1e-4)

    def test_linear_post_condition_regression(self):
        from soup_cli.utils.edit_kernels import _rank1_update

        lin = nn.Linear(16, 8, bias=False)
        key = torch.ones(16)
        delta = torch.full((8,), 0.5)
        before = lin.weight @ key
        norm = _rank1_update(lin, key, delta)
        after = lin.weight @ key
        assert norm > 0
        assert torch.allclose(after - before, delta, atol=1e-4)

    def test_conv1d_zero_key_rejected(self):
        from soup_cli.utils.edit_kernels import _rank1_update

        with pytest.raises(ValueError, match="zero norm"):
            _rank1_update(_Conv1D(16, 8), torch.zeros(16), torch.ones(8))


class TestAlphaEditProjectTransposed:
    def test_conv1d_shape_and_determinism(self):
        from soup_cli.utils.edit_kernels import _alphaedit_project

        conv = _Conv1D(16, 8)
        # Logical [out, in] update.
        upd = torch.full((8, 16), 0.3)
        p1 = _alphaedit_project(conv, upd)
        p2 = _alphaedit_project(conv, upd)
        assert tuple(p1.shape) == (8, 16)
        assert torch.allclose(p1, p2)

    def test_conv1d_projection_idempotent(self):
        from soup_cli.utils.edit_kernels import _alphaedit_project

        conv = _Conv1D(16, 8)
        upd = torch.randn(8, 16)
        once = _alphaedit_project(conv, upd)
        twice = _alphaedit_project(conv, once)
        # P is a projection: P(P(u)) == P(u).
        assert torch.allclose(once, twice, atol=1e-4)


class TestApplyKernelsGpt2EndToEnd:
    def _run(self, method, model):
        from soup_cli.utils.edit_kernels import measure_target_prob, run_edit_kernel

        tok = _FakeTok()
        before = measure_target_prob(
            model, tok, subject=_SUBJECT, target=_TARGET, device="cpu"
        )
        result = run_edit_kernel(
            model, tok, method=method, subject=_SUBJECT, target=_TARGET,
            layer=1, device="cpu",
        )
        after = measure_target_prob(
            model, tok, subject=_SUBJECT, target=_TARGET, device="cpu"
        )
        return before, result, after

    def test_rome_gpt2_changes_target(self):
        torch.manual_seed(0)
        before, result, after = self._run("rome", _FakeGPT2LM(layers=3))
        assert result.method == "rome"
        assert result.layers_edited == (1,)
        assert result.norm_delta > 0
        assert 0.0 <= after <= 1.0
        assert after > before  # the fact was edited

    def test_rome_llama_regression(self):
        torch.manual_seed(0)
        before, result, after = self._run("rome", _FakeLlamaLM(layers=3))
        assert result.norm_delta > 0
        assert after > before

    def test_memit_gpt2_runs(self):
        torch.manual_seed(1)
        _before, result, after = self._run("memit", _FakeGPT2LM(layers=3))
        assert result.method == "memit"
        assert len(result.layers_edited) >= 1
        assert result.norm_delta > 0
        assert 0.0 <= after <= 1.0

    def test_alphaedit_gpt2_runs(self):
        torch.manual_seed(2)
        _before, result, after = self._run("alphaedit", _FakeGPT2LM(layers=3))
        assert result.method == "alphaedit"
        assert result.layers_edited == (1,)
        assert result.norm_delta > 0
        assert 0.0 <= after <= 1.0

    def test_memit_gpt2_edits_full_band(self):
        """The #251 ``_proj_out_dim`` Conv1D fix makes the MEMIT band dim-check
        MATCH across uniform-width GPT-2 layers, so the whole band is edited."""
        from soup_cli.utils.edit_kernels import apply_memit_edit

        torch.manual_seed(1)
        result = apply_memit_edit(
            _FakeGPT2LM(layers=3), _FakeTok(),
            subject=_SUBJECT, target=_TARGET, layer=2, device="cpu",
        )
        # _MEMIT_BAND=3, layer=2 → band [0, 1, 2]; all same width → all edited.
        assert result.layers_edited == (0, 1, 2)

    def test_memit_raises_when_no_layer_editable(self, monkeypatch):
        """Defensive band-skip → raise branch (#251 Conv1D dim-check).

        A call-counted ``_proj_out_dim`` returns the real dim for the residual
        sizing (first call) then a mismatching dim for every band check, so the
        sole band layer is skipped and the empty-edit guard fires.
        """
        import soup_cli.utils.edit_kernels as ek

        real = ek._proj_out_dim
        state = {"n": 0}

        def fake(module):
            state["n"] += 1
            return real(module) if state["n"] == 1 else real(module) + 1

        monkeypatch.setattr(ek, "_proj_out_dim", fake)
        with pytest.raises(ValueError, match="could not edit any layer"):
            ek.apply_memit_edit(
                _FakeGPT2LM(layers=1), _FakeTok(),
                subject=_SUBJECT, target=_TARGET, layer=0, device="cpu",
            )

    def test_alphaedit_conv1d_weight_orientation(self):
        """AlphaEdit's ``.t()`` apply keeps the Conv1D weight in [in, out]
        layout — an un-transposed [out, in] apply would shape-error in add_
        (review H3)."""
        from soup_cli.utils.edit_kernels import (
            _down_proj,
            _locate_decoder_layers,
            apply_alphaedit_edit,
        )

        torch.manual_seed(3)
        model = _FakeGPT2LM(layers=3)
        down = _down_proj(_locate_decoder_layers(model), 1)
        w0 = down.weight.detach().clone()  # [in, out] = [16, 8]
        apply_alphaedit_edit(
            model, _FakeTok(),
            subject=_SUBJECT, target=_TARGET, layer=1, device="cpu",
        )
        delta_w = down.weight.detach() - w0
        assert tuple(delta_w.shape) == (16, 8)
        assert torch.isfinite(delta_w).all()
        assert float(torch.linalg.norm(delta_w)) > 0

    def test_rome_peft_wrapped_gpt2(self):
        """A PEFT-wrapped GPT-2 base is editable end-to-end through the kernel
        (review L1 — PEFT fallback at the kernel level, not just locate)."""
        from soup_cli.utils.edit_kernels import run_edit_kernel

        torch.manual_seed(0)
        model = _PeftCallable(_FakeGPT2LM(layers=3))
        result = run_edit_kernel(
            model, _FakeTok(), method="rome",
            subject=_SUBJECT, target=_TARGET, layer=1, device="cpu",
        )
        assert result.method == "rome"
        assert result.norm_delta > 0


# ===========================================================================
# #252 — atomic EditGovernor edit-count increment
# ===========================================================================


class TestGovernorAtomicIncrement:
    def test_baseline_set_on_fresh(self):
        from soup_cli.utils.edit_governor import EditGovernor

        gov = EditGovernor(base_model="m")
        assert gov._persisted_edit_count == 0

    def test_baseline_set_on_loaded(self, tmp_path, monkeypatch):
        from soup_cli.utils.edit_governor import (
            EditGovernor,
            EditGovernorStore,
            load_governor,
            save_governor,
        )

        monkeypatch.chdir(tmp_path)
        db = str(tmp_path / "g.db")
        with EditGovernorStore(db) as store:
            gov = EditGovernor(base_model="m")
            gov.record_edit(method="rome", norm_delta=0.1)
            gov.record_edit(method="rome", norm_delta=0.1)
            save_governor(store, gov)
        with EditGovernorStore(db) as store2:
            restored = load_governor(store2, "m")
        # Baseline mirrors the loaded count so a further edit merges as +1.
        assert restored.edit_count == 2
        assert restored._persisted_edit_count == 2

    def test_concurrent_save_merges_increments(self, tmp_path, monkeypatch):
        """Two governors loaded from the same state both record + save.

        With the pre-#252 absolute-write behaviour the second save would
        clobber the first to 1. The baseline-delta merge keeps both → 2.
        """
        from soup_cli.utils.edit_governor import (
            EditGovernorStore,
            load_governor,
            save_governor,
        )

        monkeypatch.chdir(tmp_path)
        db = str(tmp_path / "g.db")
        with EditGovernorStore(db) as store:
            gov_a = load_governor(store, "m")
            gov_b = load_governor(store, "m")
            gov_a.record_edit(method="rome", norm_delta=0.1)
            gov_b.record_edit(method="rome", norm_delta=0.2)
            save_governor(store, gov_a)  # persists 1
            save_governor(store, gov_b)  # MERGES → persists 2 (not clobber)
            final = load_governor(store, "m")
        assert final.edit_count == 2

    def test_concurrent_save_merges_multi_increments(self, tmp_path, monkeypatch):
        """Two governors record MULTIPLE edits each → merged is the sum of the
        deltas (3 + 2 = 5), proving the baseline-delta merge (not a naive +1
        per save) — review M4."""
        from soup_cli.utils.edit_governor import (
            EditGovernorStore,
            load_governor,
            save_governor,
        )

        monkeypatch.chdir(tmp_path)
        db = str(tmp_path / "g.db")
        with EditGovernorStore(db) as store:
            a = load_governor(store, "m")
            b = load_governor(store, "m")
            for _ in range(3):
                a.record_edit(method="rome", norm_delta=0.1)
            for _ in range(2):
                b.record_edit(method="rome", norm_delta=0.1)
            save_governor(store, a)  # persists 3
            save_governor(store, b)  # merges +2 → 5 (NOT clobber to 2, NOT +1)
            assert load_governor(store, "m").edit_count == 5

    def test_merge_onto_existing_row(self, tmp_path, monkeypatch):
        from soup_cli.utils.edit_governor import (
            EditGovernorStore,
            load_governor,
            save_governor,
        )

        monkeypatch.chdir(tmp_path)
        db = str(tmp_path / "g.db")
        with EditGovernorStore(db) as store:
            # Seed a persisted count of 5.
            g0 = load_governor(store, "m")
            for _ in range(5):
                g0.record_edit(method="rome", norm_delta=0.1)
            save_governor(store, g0)
        with EditGovernorStore(db) as store2:
            g1 = load_governor(store2, "m")  # baseline 5
            g1.record_edit(method="rome", norm_delta=0.2)  # → 6
            save_governor(store2, g1)
            final = load_governor(store2, "m")
        assert final.edit_count == 6

    def test_save_updates_in_memory_count(self, tmp_path, monkeypatch):
        """After an atomic merge, ``governor.edit_count`` reflects the merged
        value so the CLI summary shows the real count."""
        from soup_cli.utils.edit_governor import (
            EditGovernorStore,
            load_governor,
            save_governor,
        )

        monkeypatch.chdir(tmp_path)
        db = str(tmp_path / "g.db")
        with EditGovernorStore(db) as store:
            a = load_governor(store, "m")
            b = load_governor(store, "m")
            a.record_edit(method="rome", norm_delta=0.1)
            b.record_edit(method="rome", norm_delta=0.2)
            save_governor(store, a)
            save_governor(store, b)
            # b merged onto a's persisted 1 → 2; b.edit_count must reflect it.
            assert b.edit_count == 2
            # Re-saving b must NOT double-count (baseline now 2).
            save_governor(store, b)
            assert load_governor(store, "m").edit_count == 2

    def test_save_state_atomic_uses_lock(self):
        """Regression: the get+insert in save_state runs under the lock."""
        import inspect

        from soup_cli.utils.edit_governor import EditGovernorStore

        src = inspect.getsource(EditGovernorStore.save_state)
        assert "_cross_process_lock" in src
        assert "get_state" in src  # read inside the lock


# ===========================================================================
# #250 — covariance-preconditioned ROME
# ===========================================================================


class TestEstimateKeyCovariance:
    def test_shape_and_spd(self):
        from soup_cli.utils.edit_kernels import (
            _down_proj,
            _locate_decoder_layers,
            estimate_key_covariance,
        )

        model = _FakeGPT2LM()
        down = _down_proj(_locate_decoder_layers(model), 1)
        cov = estimate_key_covariance(
            model, _FakeTok(), down, _CORPUS, device="cpu",
        )
        # Covariance dim = down-proj INPUT dim (intermediate) = 16.
        assert tuple(cov.shape) == (16, 16)
        # Symmetric.
        assert torch.allclose(cov, cov.t(), atol=1e-5)
        # SPD (ridge guarantees positive eigenvalues).
        eigvals = torch.linalg.eigvalsh(cov)
        assert float(eigvals.min()) > 0.0

    def test_empty_corpus_rejected(self):
        from soup_cli.utils.edit_kernels import (
            _down_proj,
            _locate_decoder_layers,
            estimate_key_covariance,
        )

        model = _FakeGPT2LM()
        down = _down_proj(_locate_decoder_layers(model), 0)
        with pytest.raises(ValueError, match="corpus"):
            estimate_key_covariance(model, _FakeTok(), down, [], device="cpu")

    def test_bad_caps_rejected(self):
        from soup_cli.utils.edit_kernels import (
            _down_proj,
            _locate_decoder_layers,
            estimate_key_covariance,
        )

        model = _FakeGPT2LM()
        down = _down_proj(_locate_decoder_layers(model), 0)
        with pytest.raises(ValueError, match="max_prompts"):
            estimate_key_covariance(
                model, _FakeTok(), down, _CORPUS, device="cpu", max_prompts=0
            )
        with pytest.raises(ValueError, match="max_tokens"):
            estimate_key_covariance(
                model, _FakeTok(), down, _CORPUS, device="cpu", max_tokens=-1
            )
        with pytest.raises(ValueError, match="ridge"):
            estimate_key_covariance(
                model, _FakeTok(), down, _CORPUS, device="cpu", ridge=-1.0
            )

    def test_all_blank_corpus_rejected(self):
        """A corpus where every entry is blank / non-str captures no keys →
        the runtime ``count == 0`` branch raises (review H4)."""
        from soup_cli.utils.edit_kernels import (
            _down_proj,
            _locate_decoder_layers,
            estimate_key_covariance,
        )

        model = _FakeGPT2LM()
        down = _down_proj(_locate_decoder_layers(model), 0)
        with pytest.raises(ValueError, match="no key vectors"):
            estimate_key_covariance(
                model, _FakeTok(), down, ["", "   ", 123], device="cpu"
            )


class TestRank1UpdatePreconditioned:
    def test_post_condition_preserved_linear(self):
        """With C != I the ROME post-condition ``down(key*) += delta`` still
        holds exactly — the covariance only redistributes the update mass."""
        from soup_cli.utils.edit_kernels import _rank1_update

        lin = nn.Linear(16, 8, bias=False)
        key = torch.randn(16)
        delta = torch.randn(8)
        # Arbitrary SPD covariance.
        a = torch.randn(16, 16)
        cov = a @ a.t() + torch.eye(16)
        before = lin.weight @ key
        norm = _rank1_update(lin, key, delta, cov=cov)
        after = lin.weight @ key
        assert norm > 0
        assert torch.allclose(after - before, delta, atol=1e-3)

    def test_post_condition_preserved_conv1d(self):
        from soup_cli.utils.edit_kernels import _rank1_update

        conv = _Conv1D(16, 8)
        key = torch.randn(16)
        delta = torch.randn(8)
        a = torch.randn(16, 16)
        cov = a @ a.t() + torch.eye(16)
        before = key @ conv.weight
        norm = _rank1_update(conv, key, delta, cov=cov)
        after = key @ conv.weight
        assert norm > 0
        assert torch.allclose(after - before, delta, atol=1e-3)

    def test_non_finite_cov_rejected(self):
        """A non-finite covariance must raise (not silently corrupt weights).

        Review-fix: ``denom = NaN`` would slip past the bare ``<= 0`` guard.
        """
        from soup_cli.utils.edit_kernels import _rank1_update

        lin = nn.Linear(16, 8, bias=False)
        bad_cov = torch.full((16, 16), float("nan"))
        with pytest.raises(ValueError):
            _rank1_update(lin, torch.ones(16), torch.ones(8), cov=bad_cov)

    def test_singular_cov_rejected(self):
        """A singular (rank-deficient) covariance makes the solve fail /
        produce a non-finite denom → clean ValueError (review M3)."""
        from soup_cli.utils.edit_kernels import _rank1_update

        lin = nn.Linear(16, 8, bias=False)
        with pytest.raises(ValueError, match="covariance solve failed|degenerate"):
            _rank1_update(
                lin, torch.ones(16), torch.ones(8), cov=torch.zeros(16, 16)
            )

    def test_cov_changes_update_direction(self):
        """A non-identity covariance produces a different update than C=I."""
        from soup_cli.utils.edit_kernels import _rank1_update

        key = torch.randn(16)
        delta = torch.randn(8)
        a = torch.randn(16, 16)
        cov = a @ a.t() + torch.eye(16)

        lin_iden = nn.Linear(16, 8, bias=False)
        w0 = lin_iden.weight.detach().clone()
        _rank1_update(lin_iden, key, delta)
        upd_iden = lin_iden.weight.detach() - w0

        lin_cov = nn.Linear(16, 8, bias=False)
        with torch.no_grad():
            lin_cov.weight.copy_(w0)
        _rank1_update(lin_cov, key, delta, cov=cov)
        upd_cov = lin_cov.weight.detach() - w0

        assert not torch.allclose(upd_iden, upd_cov, atol=1e-3)


class TestApplyRomeWithCovCorpus:
    def test_apply_edit_rome_cov(self, monkeypatch):
        import soup_cli.utils.live_eval as live_eval
        from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan

        torch.manual_seed(0)
        model = _FakeGPT2LM(layers=3)
        tok = _FakeTok()
        monkeypatch.setattr(
            live_eval, "load_model_and_tokenizer",
            lambda *a, **k: (model, tok, "cpu"),
        )
        plan = build_edit_plan(
            base="b", method="rome", subject=_SUBJECT, target=_TARGET, layer=1,
        )
        result = apply_edit(plan, cov_corpus=_CORPUS)
        assert result.method == "rome"
        assert result.norm_delta > 0
        assert result.target_prob_after > result.target_prob_before

    def test_cov_corpus_rejected_for_memit(self, monkeypatch):
        import soup_cli.utils.live_eval as live_eval
        from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan

        # load_model_and_tokenizer must NOT be reached — the reject is before it.
        monkeypatch.setattr(
            live_eval, "load_model_and_tokenizer",
            lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")),
        )
        plan = build_edit_plan(
            base="b", method="memit", subject="s", target="t",
        )
        with pytest.raises(ValueError, match="cov-corpus"):
            apply_edit(plan, cov_corpus=_CORPUS)

    def test_cov_corpus_rejected_for_alphaedit(self, monkeypatch):
        import soup_cli.utils.live_eval as live_eval
        from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan

        monkeypatch.setattr(
            live_eval, "load_model_and_tokenizer",
            lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")),
        )
        plan = build_edit_plan(
            base="b", method="alphaedit", subject="s", target="t",
        )
        with pytest.raises(ValueError, match="cov-corpus"):
            apply_edit(plan, cov_corpus=_CORPUS)

    def test_cov_corpus_rejected_for_grace(self):
        """grace takes a different code path (codebook sidecar) but the cov
        reject still fires first — before any import / model load (review H1)."""
        from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan

        plan = build_edit_plan(base="b", method="grace", subject="s", target="t")
        with pytest.raises(ValueError, match="cov-corpus"):
            apply_edit(plan, cov_corpus=_CORPUS)

    def test_cov_reject_runs_before_governor(self, monkeypatch):
        """Order matters: a non-ROME method + cov_corpus + a governor that
        WOULD refuse must report the cov error, not the governance refusal
        (the cov check precedes governor.check_can_edit) — review M5."""
        import soup_cli.utils.live_eval as live_eval
        from soup_cli.utils.edit_governor import EditGovernor
        from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan

        monkeypatch.setattr(
            live_eval, "load_model_and_tokenizer",
            lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")),
        )
        gov = EditGovernor(base_model="b")
        gov.record_edit(method="rome", norm_delta=100.0)  # → BLOWUP: would refuse
        plan = build_edit_plan(base="b", method="memit", subject="s", target="t")
        with pytest.raises(ValueError, match="cov-corpus"):
            apply_edit(plan, governor=gov, cov_corpus=_CORPUS)


class TestLoadCovCorpus:
    def test_parses_jsonl_text_field(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        f = tmp_path / "corpus.jsonl"
        f.write_text(
            json.dumps({"text": "row one"}) + "\n"
            + json.dumps({"prompt": "row two"}) + "\n"
            + json.dumps({"content": "row three"}) + "\n",
            encoding="utf-8",
        )
        rows = _load_cov_corpus("corpus.jsonl")
        assert rows == ["row one", "row two", "row three"]

    def test_parses_raw_text_lines(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        f = tmp_path / "corpus.txt"
        f.write_text("plain line one\nplain line two\n", encoding="utf-8")
        rows = _load_cov_corpus("corpus.txt")
        assert rows == ["plain line one", "plain line two"]

    def test_outside_cwd_rejected(self):
        from soup_cli.commands.edit import _load_cov_corpus

        with pytest.raises(ValueError, match="cwd"):
            _load_cov_corpus("/etc/passwd")

    def test_null_byte_rejected(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        with pytest.raises(ValueError, match="null"):
            _load_cov_corpus("a\x00b.jsonl")

    def test_missing_file(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        with pytest.raises(FileNotFoundError):
            _load_cov_corpus("nope.jsonl")

    def test_directory_rejected(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        (tmp_path / "adir").mkdir()
        # POSIX: os.open succeeds, fstat → not S_ISREG → "regular file".
        # Windows: os.open on a dir raises PermissionError → "not readable".
        with pytest.raises(ValueError, match="regular file|not readable"):
            _load_cov_corpus("adir")

    def test_empty_corpus_rejected(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        f = tmp_path / "empty.jsonl"
        f.write_text("\n   \n", encoding="utf-8")
        with pytest.raises(ValueError, match="no usable"):
            _load_cov_corpus("empty.jsonl")

    def test_jsonl_dict_without_usable_field_skipped(self, tmp_path, monkeypatch):
        """A JSONL object with no text/prompt/content field is silently dropped
        (NOT appended as raw JSON) — review L5."""
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        f = tmp_path / "c.jsonl"
        f.write_text(
            json.dumps({"other": "x"}) + "\n" + json.dumps({"text": "y"}) + "\n",
            encoding="utf-8",
        )
        assert _load_cov_corpus("c.jsonl") == ["y"]

    def test_oversize_rejected(self, tmp_path, monkeypatch):
        """File larger than the byte cap is rejected before any read — L4."""
        import soup_cli.commands.edit as edit_mod
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        monkeypatch.setattr(edit_mod, "_MAX_COV_CORPUS_BYTES", 8)
        f = tmp_path / "big.txt"
        f.write_text("this is definitely more than eight bytes\n", encoding="utf-8")
        with pytest.raises(ValueError, match="too large"):
            _load_cov_corpus("big.txt")

    def test_line_cap_truncates(self, tmp_path, monkeypatch):
        """Reading stops at the line cap rather than consuming the whole file — L4."""
        import soup_cli.commands.edit as edit_mod
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        monkeypatch.setattr(edit_mod, "_MAX_COV_CORPUS_LINES", 2)
        f = tmp_path / "many.txt"
        f.write_text("a\nb\nc\nd\n", encoding="utf-8")
        assert _load_cov_corpus("many.txt") == ["a", "b"]

    @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink")
    def test_symlink_rejected(self, tmp_path, monkeypatch):
        from soup_cli.commands.edit import _load_cov_corpus

        monkeypatch.chdir(tmp_path)
        real = tmp_path / "real.jsonl"
        real.write_text(json.dumps({"text": "x"}) + "\n", encoding="utf-8")
        os.symlink(real, tmp_path / "link.jsonl")
        with pytest.raises(ValueError, match="symlink"):
            _load_cov_corpus("link.jsonl")


# ===========================================================================
# #147 — Mixtral in the LongLoRA allowlist
# ===========================================================================


class TestIsMixtralModel:
    def test_basic(self):
        from soup_cli.utils.longlora import is_mixtral_model

        assert is_mixtral_model("mistralai/Mixtral-8x7B-v0.1") is True
        assert is_mixtral_model("mistralai/Mixtral-8x22B-Instruct-v0.1") is True

    def test_bare_token(self):
        """The start-of-string anchor matches a lone ``mixtral-...`` id (M1)."""
        from soup_cli.utils.longlora import is_mixtral_model

        assert is_mixtral_model("Mixtral-8x7B-v0.1") is True

    def test_not_plain_mistral(self):
        from soup_cli.utils.longlora import is_mixtral_model

        assert is_mixtral_model("mistralai/Mistral-7B-v0.1") is False

    def test_word_boundary(self):
        from soup_cli.utils.longlora import is_mixtral_model

        assert is_mixtral_model("my-mixtralish-finetune") is False
        assert is_mixtral_model("unmixtral-7b") is False

    def test_input_guards(self):
        from soup_cli.utils.longlora import is_mixtral_model

        assert is_mixtral_model("") is False
        with pytest.raises(TypeError):
            is_mixtral_model(None)  # type: ignore[arg-type]
        with pytest.raises(ValueError):
            is_mixtral_model("a\x00b")
        assert is_mixtral_model("a" * 1024) is False


class TestMixtralStillNotMistral:
    def test_is_mistral_model_excludes_mixtral(self):
        from soup_cli.utils.longlora import is_mistral_model

        # Regression — is_mistral_model stays narrow; Mixtral is detected by
        # the dedicated is_mixtral_model helper.
        assert is_mistral_model("mistralai/Mixtral-8x7B-v0.1") is False


class TestMixtralInAllowlist:
    def test_supported(self):
        from soup_cli.utils.longlora import is_supported_longlora_arch

        assert is_supported_longlora_arch("mistralai/Mixtral-8x7B-v0.1") is True
        assert (
            is_supported_longlora_arch("mistralai/Mixtral-8x22B-Instruct-v0.1")
            is True
        )

    def test_unsupported_unchanged(self):
        from soup_cli.utils.longlora import is_supported_longlora_arch

        assert is_supported_longlora_arch("google/gemma-2-9b") is False
        assert is_supported_longlora_arch("databricks/dbrx-base") is False

    def test_separate_qkv_families_includes_mixtral(self):
        from soup_cli.utils.longlora import _SEPARATE_QKV_FAMILIES

        assert "Mixtral" in _SEPARATE_QKV_FAMILIES

    def test_defensive_surface(self):
        """The rewired ``or``-chain must still swallow non-str / null-byte input
        (returns False, never raises) — review M2."""
        from soup_cli.utils.longlora import is_supported_longlora_arch

        assert is_supported_longlora_arch(None) is False
        assert is_supported_longlora_arch(123) is False
        assert is_supported_longlora_arch("a\x00b") is False


class TestValidateLongloraCompatMixtral:
    def test_accepts_mixtral(self, monkeypatch):
        from soup_cli.utils import longlora

        monkeypatch.setattr(
            "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False
        )
        # Should not raise.
        longlora.validate_longlora_compat(
            model_name="mistralai/Mixtral-8x7B-v0.1",
            task="sft",
            backend="transformers",
            use_ring_attention=False,
        )

    def test_error_message_lists_mixtral(self, monkeypatch):
        from soup_cli.utils import longlora

        monkeypatch.setattr(
            "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False
        )
        with pytest.raises(ValueError) as exc:
            longlora.validate_longlora_compat(
                model_name="google/gemma-2-9b",
                task="sft",
                backend="transformers",
                use_ring_attention=False,
            )
        assert "Mixtral" in str(exc.value)

    def test_mixtral_ring_attention_rejected(self, monkeypatch):
        """Now that Mixtral passes the arch gate, the downstream ring-attention
        exclusivity becomes reachable for it — confirm it still fires (M6)."""
        from soup_cli.utils import longlora

        monkeypatch.setattr(
            "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False
        )
        with pytest.raises(ValueError, match="ring"):
            longlora.validate_longlora_compat(
                model_name="mistralai/Mixtral-8x7B-v0.1",
                task="sft",
                backend="transformers",
                use_ring_attention=True,
            )


class TestMixtralForwardOverride:
    def _fake_attn_model(self, cls_name):
        # The override matches on the attention module's CLASS NAME, so the
        # instance's class must literally be named e.g. ``MixtralAttention``.
        def _init(self):
            nn.Module.__init__(self)
            self.head_dim = 4
            self.num_heads = 2
            self.num_key_value_heads = 2
            self.q_proj = nn.Linear(8, 8, bias=False)
            self.k_proj = nn.Linear(8, 8, bias=False)
            self.v_proj = nn.Linear(8, 8, bias=False)

        attn_cls = type(cls_name, (nn.Module,), {"__init__": _init})
        model = nn.Module()
        model.attn = attn_cls()  # registers as a submodule
        return model

    def test_patches_mixtral_attention(self):
        from soup_cli.utils.longlora import LongLoRAForwardOverride

        model = self._fake_attn_model("MixtralAttention")
        with LongLoRAForwardOverride(model, group_size=4):
            assert getattr(
                model.attn.q_proj.forward, "_soup_longlora_patched", False
            )
            assert getattr(
                model.attn.k_proj.forward, "_soup_longlora_patched", False
            )
        # Restored on exit.
        assert not getattr(
            model.attn.q_proj.forward, "_soup_longlora_patched", False
        )

    def test_llama_attention_still_patched(self):
        from soup_cli.utils.longlora import LongLoRAForwardOverride

        model = self._fake_attn_model("LlamaAttention")
        with LongLoRAForwardOverride(model, group_size=4):
            assert getattr(
                model.attn.q_proj.forward, "_soup_longlora_patched", False
            )


class TestMixtralSchemaGate:
    def test_schema_accepts_mixtral(self, monkeypatch):
        from soup_cli.config.loader import load_config_from_string

        monkeypatch.setattr(
            "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False
        )
        cfg = load_config_from_string(
            "base: mistralai/Mixtral-8x7B-v0.1\n"
            "task: sft\n"
            "data:\n"
            "  train: data.jsonl\n"
            "training:\n"
            "  use_longlora: true\n"
        )
        assert cfg.training.use_longlora is True


# ===========================================================================
# CLI plumbing
# ===========================================================================


class TestCliCovCorpus:
    def test_edit_set_help_has_cov_corpus(self):
        import re

        from typer.testing import CliRunner

        from soup_cli.cli import app

        result = CliRunner().invoke(app, ["edit", "set", "--help"])
        assert result.exit_code == 0, result.output
        clean = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
        assert "cov-corpus" in clean


# ===========================================================================
# Patch invariants
# ===========================================================================


class TestPatchInvariants:
    def test_version_bumped(self):
        import soup_cli

        parts = soup_cli.__version__.split(".")
        assert (int(parts[0]), int(parts[1]), int(parts[2])) >= (0, 71, 16)

    @pytest.mark.parametrize(
        "module",
        [
            "soup_cli.utils.edit_kernels",
            "soup_cli.utils.edit_governor",
            "soup_cli.utils.longlora",
            "soup_cli.commands.edit",
        ],
    )
    def test_no_top_level_torch(self, module):
        import importlib

        mod = importlib.import_module(module)
        with open(mod.__file__, encoding="utf-8") as fh:
            src = fh.read()
        for line in src.splitlines():
            assert not line.startswith("import torch"), module
            assert not line.startswith("from torch"), module

    def test_edit_kernels_no_top_level_safetensors(self):
        import importlib

        mod = importlib.import_module("soup_cli.utils.edit_kernels")
        with open(mod.__file__, encoding="utf-8") as fh:
            src = fh.read()
        assert "estimate_key_covariance" in src