"""Filler augmentation for long-context anchor training (SPEC.md §12.4), on a tiny random
Qwen3.5 (CPU): a short record embedded in unrelated filler text keeps every gold span on
its own text and every target of the plain example; negative (filler only) examples
train top-level anchored arrays as empty; the corpus uses train-split documents only
and refuses ExtractBench; micro-batches materialize filler draws before packing (the
memory planner sees real examples) within the context limit.
"""

from __future__ import annotations

import json
import random
import zlib
from pathlib import Path

import pytest

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

from test_anchors import NESTED, TEAM_CONTEXT  # noqa: E402
from test_array_decode import _team_record  # noqa: E402
from test_model import SNAPSHOT, tiny_lm  # noqa: E402

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

FILLER_WORDS = ("The quarterly filing discusses liquidity, covenants and segment revenue "
                "across fiscal periods; management notes risks in supply and demand. ")


@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))
    torch.manual_seed(1)
    with torch.no_grad():
        for p in m.heads.parameters():
            p.normal_(0.0, 0.5)
    return m


def _val(doc_id: str) -> bool:
    from deg.datagen.v3 import VAL_MODULUS

    return zlib.crc32(doc_id.encode()) % VAL_MODULUS == 0


@pytest.fixture()
def corpus(tmp_path: Path):
    from deg.train import FillerCorpus

    (tmp_path / "text").mkdir()
    rows = []
    val_id = next(f"doc-{i}" for i in range(10_000) if _val(f"doc-{i}"))
    for i, doc_id in enumerate(["doc-a", "doc-b", val_id]):
        body = (f"[{doc_id}] " if doc_id == val_id else "") + FILLER_WORDS * (40 + 10 * i)
        tp = tmp_path / "text" / f"{doc_id}.txt"
        tp.write_text(body)
        rows.append({"id": doc_id, "text_path": str(tp)})
    (tmp_path / "docs.jsonl").write_text("".join(json.dumps(r) + "\n" for r in rows))
    c = FillerCorpus(tmp_path)
    assert {d for d, _ in c.docs} == {"doc-a", "doc-b"}  # the v3 val doc is never filler
    return c


def _targets(ex, text: str):
    """Every training target as (kind, ptr, gold text), independent of token positions."""
    def at(i: int, j: int) -> str:
        return text[ex.ctx.offsets[i][0]:ex.ctx.offsets[j][1]]

    out = []
    for it in ex.items:
        if it.kind == "anchor":
            out.append(("anchor", it.ptr, tuple(at(t, t) for t in it.anchors)))
        elif it.kind == "extract":
            out.append(("extract", it.ptr, tuple(
                at(a, b) if a >= 0 else None for a, b in zip(it.starts, it.ends))))
        else:
            out.append((it.kind, it.ptr, tuple(it.target or ()), it.aux))
    return sorted(out, key=repr)


# ----------------------------------------------------------------------- filler


def test_corpus_refuses_extractbench(tmp_path):
    from deg.train import FillerCorpus

    d = tmp_path / "ExtractBench"
    d.mkdir()
    (d / "docs.jsonl").write_text("")
    with pytest.raises(ValueError):
        FillerCorpus(d)


def test_filler_record_keeps_every_target_on_its_text(model, corpus):
    from deg.train import deg_example, filler_record, record_text

    rec = _team_record()
    plain = deg_example(model, rec, random.Random(0), shuffle_p=0.0, anchor_count=True)
    want = _targets(plain, record_text(rec))
    assert any(t[0] == "anchor" for t in want) and any(t[0] == "count" for t in want)
    seen_offsets = set()
    for seed in range(6):
        rng = random.Random(seed)
        aug = filler_record(model, rec, corpus, rng, 200)
        text = aug["context"]
        assert TEAM_CONTEXT in text and len(text) > len(TEAM_CONTEXT) + 200
        shift = text.index(TEAM_CONTEXT)
        seen_offsets.add(shift)
        for k, (a, b) in aug["spans"].items():
            ra, rb = rec["spans"][k]
            assert (a - shift, b - shift) == (ra, rb) and text[a:b] == TEAM_CONTEXT[ra:rb]
        ex = deg_example(model, aug, random.Random(0), shuffle_p=0.0, anchor_count=True)
        assert len(ex.ctx.ids) > len(plain.ctx.ids) + 150
        assert _targets(ex, text) == want  # same targets, same gold text
    assert len(seen_offsets) > 2  # random placement, including the middle


def test_negative_example_trains_top_level_anchors_empty(model, corpus):
    from deg.train import negative_example

    ex = negative_example(model, _team_record(), corpus, random.Random(0), 150,
                          anchor_count=True)
    assert ex is not None and ex.aug == "negative"
    assert 140 <= len(ex.ctx.offsets) <= 151
    anchors = [it for it in ex.items if it.kind == "anchor"]
    counts = [it for it in ex.items if it.kind == "count"]
    assert [it.ptr for it in anchors] == ["/teams"] and anchors[0].anchors == []
    assert len(counts) == 1 and counts[0].aux and counts[0].target.index(1.0) == 0
    assert not any(it.kind in ("extract", "choice", "score", "noul") for it in ex.items)
    min1 = {**NESTED, "properties": {"teams": {**NESTED["properties"]["teams"],
                                               "minItems": 1}}}
    assert negative_example(model, {**_team_record(), "schema": min1}, corpus,
                            random.Random(0), 150) is None


def test_filler_draws_materialize_before_packing_within_the_limit(model, corpus):
    from deg.train import (
        FillerDraw,
        Source,
        batches,
        deg_example,
        filler_example,
    )

    rec = _team_record()
    ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0, anchor_count=True)
    ex.rec = rec
    src = Source(Path("x.jsonl"))
    src.pool = [ex]
    max_ctx = 300
    seen: list[list] = []

    def fits(batch):
        seen.append(list(batch))
        return True

    def augment(e, r):
        return FillerDraw(e, r.randint(len(e.ctx.ids) + 64, max_ctx), r.random() < 0.3)

    def materialize(d, r):
        return filler_example(model, d, corpus, r, max_ctx, anchor_count=True)

    gen = batches([src], None, random.Random(0), 10_000, 0.0, 4, draw=8,
                  materialize=materialize, fits=fits, augment=augment)
    kinds = set()
    for _ in range(6):
        for e in next(gen):
            assert len(e.ctx.ids) <= max_ctx
            kinds.add(e.aug)
    assert kinds == {"filler", "negative"}
    assert seen and all(not isinstance(e, FillerDraw) for b in seen for e in b)


def test_filler_example_falls_back_to_the_plain_example(model, corpus):
    from deg.train import FillerDraw, deg_example, filler_example

    rec = _team_record()
    ex = deg_example(model, rec, random.Random(0))
    ex.rec = rec
    # no room for filler under the context limit -> the plain example
    got = filler_example(model, FillerDraw(ex, len(ex.ctx.ids) + 500), corpus,
                         random.Random(0), len(ex.ctx.ids) + 3)
    assert got is ex