"""Training trunks longer than DENSE_PREFILL_MAX (SPEC.md §12.4): the chunked read is
gradient-checkpointed per layer and per chunk, and equals the dense (one-shot,
checkpointed) path in hidden states, continued state, fork outputs and gradients, on a
tiny random Qwen3.5 (CPU)."""

from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("transformers")

from test_model import CONTEXT, SNAPSHOT, _LargestTensor, tiny_lm  # noqa: E402

from deg.model import load_model, load_tokenizer  # noqa: E402


@pytest.fixture(scope="module")
def model():
    if SNAPSHOT is None:
        pytest.skip("Qwen3.5-0.8B tokenizer not cached")
    m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(),
                   tokenizer=load_tokenizer(SNAPSHOT))
    m.backbone.grad_ckpt = True
    m.backbone.text.embed_tokens.weight.requires_grad_(False)
    return m


def _run(model, rows: list[list[int]], forks: list[list[int]], fork_rows: list[int]):
    bb = model.backbone
    for p in bb.parameters():
        p.grad = None
    width = max(len(r) for r in rows)
    ids = torch.full((len(rows), width), model.renderer.pad_id, dtype=torch.long)
    for i, r in enumerate(rows):
        ids[i, :len(r)] = torch.tensor(r)
    lengths = torch.tensor([len(r) for r in rows])
    bb.train()
    try:
        h, state = bb(ids, lengths)
        fh = model.run_forks(state, torch.tensor(fork_rows), [tuple(f) for f in forks])
        loss = sum((h[i, :len(r)].float() ** 2).mean() for i, r in enumerate(rows))
        loss = loss + sum((x.float() ** 2).mean() for x in fh)
        loss.backward()
    finally:
        bb.eval()
    grads = {n: p.grad.clone() for n, p in bb.named_parameters() if p.grad is not None}
    return h.detach(), state, [x.detach() for x in fh], grads


def test_chunked_checkpointed_trunk_equals_dense(model, monkeypatch):
    import deg.model as dm

    rend = model.renderer
    a = rend.context(" ".join([CONTEXT] * 3)).ids
    b = rend.context(" ".join([CONTEXT] * 2)).ids
    q = rend.enc("FIELD /x\nKIND: copy\nANSWER:")
    assert len(a) > 100
    dense = _run(model, [a, b], [q, q[:5], q], [0, 1, 1])
    monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 48)
    monkeypatch.setattr(dm, "PREFILL_CHUNK", 32)
    calls = []
    orig = dm.Backbone._forward_chunked_ckpt

    def spy(self, *args, **kw):
        calls.append(1)
        return orig(self, *args, **kw)

    monkeypatch.setattr(dm.Backbone, "_forward_chunked_ckpt", spy)
    with _LargestTensor() as big:
        chunked = _run(model, [a, b], [q, q[:5], q], [0, 1, 1])
    assert calls, "long training trunk did not take the checkpointed chunked path"
    L = len(a)
    assert big.largest < 2 * L * L // 3, big.largest  # no [B, L, L] mask
    (h1, s1, f1, g1), (h2, s2, f2, g2) = dense, chunked
    for i, n in enumerate((len(a), len(b))):
        torch.testing.assert_close(h2[i, :n], h1[i, :n], atol=2e-4, rtol=2e-4)
    for (r, t), (r1, t1) in zip(s2.lin, s1.lin):
        torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4)
        torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4)
    for (k, v), (k1, v1) in zip(s2.kv, s1.kv):
        assert k.shape == k1.shape
        m = s1.kv_mask[:, None, :, None]
        torch.testing.assert_close(k * m, k1 * m, atol=2e-4, rtol=2e-4)
        torch.testing.assert_close(v * m, v1 * m, atol=2e-4, rtol=2e-4)
    assert torch.equal(s2.kv_mask, s1.kv_mask) and torch.equal(s2.lengths, s1.lengths)
    for x, y in zip(f2, f1):
        torch.testing.assert_close(x, y, atol=2e-4, rtol=2e-4)
    assert g1.keys() == g2.keys() and g1
    for n in g1:
        torch.testing.assert_close(g2[n], g1[n], atol=1e-4, rtol=1e-3, msg=n)


def test_inference_long_read_unchanged_by_training_path(model, monkeypatch):
    """Without autograd the chunked inference read (``_forward_chunked``) still runs."""
    import deg.model as dm

    monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 48)
    monkeypatch.setattr(dm, "PREFILL_CHUNK", 32)
    called = []
    monkeypatch.setattr(dm.Backbone, "_forward_chunked_ckpt",
                        lambda *a, **k: called.append(1))
    ids = model.renderer.context(" ".join([CONTEXT] * 2)).ids
    with torch.no_grad():
        model.backbone(torch.tensor([ids]), torch.tensor([len(ids)]))
    assert not called