test_distributed_knowledge_transfer.py

"""Multi-instance proof: closed-loop RSI telemetry -> aggregator -> BEP-46 update -> receiver hot-reload / rollback.

Every instance is a real `ArmigerUltraPipeline` loaded from this directory (fp32 ONNX backbone, the full
929,408-key zero-cue catalog, `swarm_adapter.pt`). Nothing is mocked: telemetry is hybrid X25519 + ML-KEM-768
encrypted, manifests carry real Ed25519 + WOTS+-SHAKE256 signatures, and the receiver's decisions come from its
unchanged `classify` path.

Set ARMIGER_PROOF_RESULTS=/path/results.json to dump the measured numbers used in
PROOF_OF_DISTRIBUTED_KNOWLEDGE_TRANSFER.md.
"""
from __future__ import annotations

import dataclasses
import hashlib
import json
import math
import os
from pathlib import Path
import random
import sys
import time
from typing import Any, Dict, List, Sequence

import numpy as np
import pytest

_THIS_DIR = Path(__file__).resolve().parent
if str(_THIS_DIR) not in sys.path:
    sys.path.insert(0, str(_THIS_DIR))

from armiger_ultra_pipeline import (  # noqa: E402
    ArmigerUltraPipeline,
    PeerAttestation,
    RSIUpdateAggregator,
    WOTSPlusSHAKE256,
    encode_safetensors,
    generate_federation_signer,
)

EPOCH = "2026-10"
RESULTS: Dict[str, Any] = {}
PEER_ASNS = [64512, 64513, 64514, 64515, 64516, 64517]  # RFC 6996 private-use ASNs (synthetic peers)

# Release-name variants: contributors each grabbed a different release; the receiver sees a name nobody contributed.
CONTRIB_VARIANTS = [
    "{T} ({Y}) [1080p]",
    "{D}.{Y}.720p.WEB-DL.H264",
    "{T} {Y} 1080p WEBRip x265",
    "{D}.{Y}.2160p.WEB.H265",
    "{T} {Y} 1080p WEB-DL",
    "{D}.{Y}.1080p.WEB.h264",
]
RECEIVER_VARIANT = "{T} ({Y}) [720p] [WEB-DL]"
CONTROL_RECEIVER_VARIANT = "{D}.{Y}.720p.WEB-DL.H264"
TARGETS = ["Hollow Meridian", "The Quiet Orchard", "Low Tide at Carrow", "Northern Relay", "Kettle Point"]
CONTROL = "Ninefold"   # verified TV, but the receiver's backbone is confidently (>=0.88) wrong on its release name
DECOY = "Saltmarsh"    # seen by one contributor only -> must not reach the update


@pytest.fixture(scope="module", autouse=True)
def _dump_results():
    yield
    out = os.environ.get("ARMIGER_PROOF_RESULTS")
    if out:
        Path(out).write_text(json.dumps(RESULTS, indent=2, sort_keys=True, default=str))


@pytest.fixture(scope="module")
def federation() -> Dict[str, Any]:
    return {
        "signers": [generate_federation_signer() for _ in range(3)],  # 2-of-3 threshold
        "peers": [{"asn": asn, "signer": generate_federation_signer()} for asn in PEER_ASNS],
    }


def _configure(pipe: ArmigerUltraPipeline, fed: Dict[str, Any]) -> ArmigerUltraPipeline:
    pipe.configure_dht_update_trust(
        trusted_signers={s[1]: f"federation-{i}" for i, s in enumerate(fed["signers"])},
        trusted_pq_signers={s[3]: f"federation-{i}" for i, s in enumerate(fed["signers"])},
        threshold_m=2,
        trusted_attestors={p["signer"][1]: {"asn": p["asn"], "pq_pubkey_hex": p["signer"][3]} for p in fed["peers"]},
        min_distinct_asns=5,
        require_pq=True,
    )
    return pipe


@pytest.fixture()
def make_instance(federation):
    created: List[ArmigerUltraPipeline] = []

    def _make() -> ArmigerUltraPipeline:
        pipe = _configure(ArmigerUltraPipeline.from_Directory(_THIS_DIR, preload_catalog=True), federation)
        created.append(pipe)
        return pipe

    yield _make
    for pipe in created:
        pipe.close()


# ---------------------------------------------------------------------------------------------
# Physical-import fixtures (*Arr webhooks, qBittorrent completion, ffprobe stream inspection)
# ---------------------------------------------------------------------------------------------

def _name(title: str, variant: str, year: int = 2026) -> str:
    return variant.format(T=title, D=title.replace(" ", "."), Y=year)


def _tv_pack_files(title: str, year: int = 2026) -> List[Dict[str, Any]]:
    return [{"path": f"{title} ({year})/0{j}.mkv", "size": 2_300_000_000} for j in (1, 2, 3)]


def _ffprobe(size: int, duration: float = 3120.0) -> Dict[str, Any]:
    return {
        "streams": [
            {"codec_type": "video", "codec_name": "h264", "width": 1920, "height": 1080},
            {"codec_type": "audio", "codec_name": "eac3", "channels": 6, "tags": {"language": "eng"}},
            {"codec_type": "subtitle", "codec_name": "subrip", "tags": {"language": "eng"}},
        ],
        "format": {"format_name": "matroska,webm", "duration": str(duration), "size": str(size)},
    }


def _series_webhook(series_title: str, catalog_id: int, release_title: str, files: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    size = int(sum(f["size"] for f in files))
    return {
        "eventType": "Download",
        "series": {"title": series_title, "seriesCatalogId": catalog_id},
        "episodes": [{"seasonNumber": 1, "episodeNumber": j + 1} for j in range(len(files))],
        "release": {"releaseTitle": release_title, "size": size},
        "episodeFile": {
            "relativePath": files[0]["path"], "size": size, "duration": 3120,
            "mediaInfo": {"videoCodec": "h264", "audioFormat": "eac3", "audioChannels": 5.1,
                          "audioLanguages": "eng", "subtitles": "eng"},
        },
    }


def _movie_webhook(title: str, year: int, catalog_id: int, release_title: str, files: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    size = int(sum(f["size"] for f in files))
    return {
        "eventType": "Download",
        "movie": {"title": title, "year": year, "movieCatalogId": catalog_id},
        "release": {"releaseTitle": release_title, "size": size},
        "movieFile": {
            "relativePath": files[0]["path"], "size": size, "duration": 6400,
            "mediaInfo": {"videoCodec": "h264", "audioFormat": "eac3", "audioChannels": 5.1,
                          "audioLanguages": "eng", "subtitles": "eng"},
        },
    }


def _qbit_completion(name: str, size: int) -> Dict[str, Any]:
    return {"hash": hashlib.sha1(name.encode()).hexdigest(), "name": name, "size": size,
            "progress": 1.0, "state": "stalledUP", "content_path": f"/downloads/complete/{name}"}


def _export(pipe, agg, payload, *, adapter, row, ffprobe=None):
    ok, reason, bundle = pipe.export_rsi_telemetry_bundle(
        payload, adapter=adapter, torrent_row=row, ffprobe_output=ffprobe, epoch=EPOCH,
        aggregator_x25519_pub=agg.x25519_pub, aggregator_mlkem_ek=agg.mlkem_ek,
    )
    assert ok, reason
    assert bundle.is_post_quantum and len(bytes.fromhex(bundle.ephemeral_pubkey_hex)) == 32 + 1088
    assert bundle.cipher_suite == "X25519-MLKEM768-ChaCha20Poly1305-SHA3"
    return bundle


def _attest_all(peers, fed, payload):
    attestations, reviews = [], []
    for pipe, peer in zip(peers, fed["peers"]):
        att, review = pipe.attest_dht_update(payload, signer=peer["signer"], asn=peer["asn"], peer_id=f"peer-asn{peer['asn']}")
        reviews.append(review)
        if att is not None:
            attestations.append(att)
    return attestations, reviews


def _forged_attestations(fed, payload, peers=None) -> List[PeerAttestation]:
    """Attacker model for Test 3: the attestor KEYS are compromised, so attestations are signed without review."""
    target = hashlib.sha256(payload).hexdigest()
    out = []
    for peer in (peers or fed["peers"]):
        ed_priv, ed_pub, pq_priv, pq_pub = peer["signer"]
        unsigned = PeerAttestation(peer_id=f"asn{peer['asn']}", asn=peer["asn"], peer_pubkey_hex=ed_pub,
                                   timestamp=int(time.time()), signature_hex="")
        msg = unsigned.canonical_bytes(target)
        out.append(dataclasses.replace(unsigned, signature_hex=ed_priv.sign(msg).hex(), pq_pubkey_hex=pq_pub,
                                       pq_signature_hex=WOTSPlusSHAKE256.sign(pq_priv, msg).hex()))
    return out


def _cat_probs(pred) -> Dict[str, float]:
    p = pred.heads["category"]["probs"]
    return {"movie": round(p["movie"], 4), "tv_show": round(p["tv_show"], 4)}


# =============================================================================================
# Test 1: Instance A learns zero-cue titles -> DHT -> Instance B improves accuracy
# =============================================================================================

def test_1_zero_cue_knowledge_transfers_from_instance_a_to_instance_b(make_instance, federation):
    instance_b = make_instance()
    b_rows = {t: {"name": _name(t, RECEIVER_VARIANT), "size": 6_900_000_000, "files": _tv_pack_files(t)} for t in TARGETS}
    control_row = {"name": _name(CONTROL, CONTROL_RECEIVER_VARIANT), "size": 6_900_000_000, "files": _tv_pack_files(CONTROL)}

    # --- Before: B cannot disambiguate the brand-new bare-title 2026 releases --------------------------
    before = {t: instance_b.classify(r) for t, r in b_rows.items()}
    for t, pred in before.items():
        src = pred.plugin_trace["plugin_3_zero_cue_catalog"]["resolution_source"]
        assert pred.category != "tv_show", f"precondition: B must start wrong on {t}"
        assert src == "neural_backbone", f"precondition: {t} must be absent from the shipped catalog"
        assert max(_cat_probs(pred).values()) < 0.88, f"precondition: {t} must be inside the catalog gate"
    control_before = instance_b.classify(control_row)
    assert control_before.category == "movie" and max(_cat_probs(control_before).values()) >= 0.88

    # --- Instance A + 5 peers (distinct ASNs) import the releases through SeriesWebhookAdapter + ffprobe ------------
    agg = RSIUpdateAggregator(epoch=EPOCH)
    instance_a = make_instance()
    peers = [instance_a] + [make_instance() for _ in range(5)]
    bundles = []
    for i, pipe in enumerate(peers):
        for k, title in enumerate(TARGETS + [CONTROL]):
            name = _name(title, CONTRIB_VARIANTS[i])
            files = _tv_pack_files(title)
            row = {"name": name, "size": sum(f["size"] for f in files), "files": files}
            bundles.append(_export(pipe, agg, _series_webhook(title, 910_001 + k, name, files),
                                   adapter="series_webhook", row=row, ffprobe=_ffprobe(row["size"])))
    decoy_files = _tv_pack_files(DECOY)
    decoy_name = _name(DECOY, CONTRIB_VARIANTS[0])
    bundles.append(_export(instance_a, agg, _series_webhook(DECOY, 919_999, decoy_name, decoy_files),
                           adapter="series_webhook", row={"name": decoy_name, "size": 6_900_000_000, "files": decoy_files},
                           ffprobe=_ffprobe(6_900_000_000)))
    # An unplayable "import" is refused before any telemetry exists
    bad_ok, bad_reason, bad_bundle = instance_a.export_rsi_telemetry_bundle(
        _series_webhook("Kettle Point", 910_005, _name("Kettle Point", CONTRIB_VARIANTS[0]), _tv_pack_files("Kettle Point")),
        adapter="series_webhook", ffprobe_output={"streams": [], "format": {}}, epoch=EPOCH,
        aggregator_x25519_pub=agg.x25519_pub, aggregator_mlkem_ek=agg.mlkem_ek)
    assert not bad_ok and bad_bundle is None and "invalid_unplayable_media_stream" in bad_reason

    # --- Aggregator: decrypt/authenticate, refuse replay + tamper, unanimous multi-contributor consensus ---
    tampered = dataclasses.replace(bundles[0], ciphertext_hex=bundles[0].ciphertext_hex[:-2] + ("00" if bundles[0].ciphertext_hex[-2:] != "00" else "01"))
    first = agg.ingest_bundles([tampered] + bundles)
    replay = agg.ingest_bundles([bundles[3]])
    assert first == {"accepted": len(bundles), "rejected": 1} and replay == {"accepted": 0, "rejected": 1}
    assert any(r.startswith("aead_or_json_rejected") for r in agg.rejected) and "duplicate_bundle_replay" in agg.rejected
    plaintext = json.dumps(agg.observations).lower()
    for title in TARGETS + [CONTROL, DECOY]:
        assert title.lower() not in plaintext and title.lower().replace(" ", ".") not in plaintext
    delta, zc_stats = agg.build_zero_cue_delta(min_support=5)
    assert zc_stats == {"keys_seen": 7, "accepted": 6, "conflicting": 0, "insufficient_support": 1}

    payload = agg.build_update_payload(zero_cue_catalog_delta=delta)
    attestations, reviews = _attest_all(peers, federation, payload)
    assert [r.status for r in reviews] == ["review_passed"] * 6
    assert len(attestations) == 6
    manifest = RSIUpdateAggregator.sign_manifest(payload, seq=1, signers=federation["signers"][:2],
                                                 version="zero-cue-2026.10.18", min_asn_quorum=5)

    # --- Instance B verifies, stages, canary-gates and hot-reloads -----------------------------------
    fp_before = instance_b.dht_state_fingerprint()
    result = instance_b.stage_and_apply_dht_update(manifest, payload, attestations=attestations, expected_prev_seq=0)
    assert result.applied and result.status == "committed", result.reason
    assert (result.seq_before, result.seq_after) == (0, 1)
    assert all(s["ok"] for s in result.verification_trace), result.verification_trace
    assert result.regressions == [] and result.state_fingerprint_after != fp_before

    after = {t: instance_b.classify(r) for t, r in b_rows.items()}
    correct = sum(p.category == "tv_show" for p in after.values())
    assert correct == len(TARGETS)  # 5/5 = 100%
    for pred in after.values():
        assert pred.plugin_trace["plugin_3_zero_cue_catalog"]["resolution_source"].startswith("zero_cue_")
    control_after = instance_b.classify(control_row)
    assert control_after.category == "movie"  # confident backbone predictions are never overridden by catalog data
    assert control_after.plugin_trace["plugin_3_zero_cue_catalog"]["resolution_source"] == "neural_backbone"

    # --- Replay of the same signed update is refused and changes nothing ----------------------------
    fp_committed = instance_b.dht_state_fingerprint()
    replayed = instance_b.stage_and_apply_dht_update(manifest, payload, attestations=attestations)
    assert replayed.status == "rejected_verification" and "stale_or_replayed_seq" in replayed.reason
    assert instance_b.dht_state_fingerprint() == fp_committed

    RESULTS["test_1"] = {
        "receiver_release_names": {t: r["name"] for t, r in b_rows.items()},
        "before": {t: {"category": p.category, **_cat_probs(p)} for t, p in before.items()},
        "after": {t: {"category": p.category, **_cat_probs(p),
                      "source": p.plugin_trace["plugin_3_zero_cue_catalog"]["resolution_source"]} for t, p in after.items()},
        "accuracy_before": f"{sum(p.category == 'tv_show' for p in before.values())}/5",
        "accuracy_after": f"{correct}/5",
        "control": {"name": control_row["name"], "before": {"category": control_before.category, **_cat_probs(control_before)},
                    "after": {"category": control_after.category, **_cat_probs(control_after)}},
        "bundles_exported": len(bundles), "aggregator_rejections": list(agg.rejected), "zero_cue_consensus": zc_stats,
        "payload_sha256": hashlib.sha256(payload).hexdigest(), "payload_bytes": len(payload),
        "manifest": {"seq": manifest.seq, "infohash": manifest.infohash, "ed25519_signers": len(manifest.signatures),
                     "wots_signers": len(manifest.pq_signatures),
                     "wots_sig_bytes": len(bytes.fromhex(next(iter(manifest.pq_signatures.values()))))},
        "peer_reviews": [{"status": r.status, "elapsed_ms": r.elapsed_ms,
                          "ledger": next((s["detail"] for s in r.verification_trace if s["step"] == "ledger_confirmation"), None)} for r in reviews],
        "receiver_trace": result.verification_trace,
        "receiver_canary_baseline": result.baseline_canary, "receiver_canary_candidate": result.candidate_canary,
        "receiver_elapsed_ms": result.elapsed_ms,
        "replay": {"status": replayed.status, "reason": replayed.reason},
    }


# =============================================================================================
# Test 2: distributed calibration + SwarmResidualAdapter weight-delta transfer
# =============================================================================================

_A = ["Amber", "Cobalt", "Willow", "Granite", "Thistle", "Harrow", "Juniper", "Marsh", "Copper", "Linden", "Briar", "Oriel"]
_B = ["Reach", "Hollow", "Ferry", "Station", "Lantern", "Crossing", "Ledger", "Acre", "Signal", "Tide", "Ridge", "Orchard"]
_G = ["NTb", "FLUX", "SiGMA", "EDITH", "CMRG"]


def _clean_swarm(rng: random.Random, i: int) -> Dict[str, Any]:
    if i % 4 == 0:  # legitimately dead swarm with real completion history
        return {"tracker_known": True, "tracker_seeders": 0, "tracker_leechers": rng.randint(10, 40),
                "tracker_completed": rng.randint(100, 3000), "prev_seeders": rng.randint(0, 9)}
    return {"tracker_known": True, "tracker_seeders": rng.randint(3, 500), "tracker_leechers": rng.randint(0, 80),
            "tracker_completed": rng.randint(20, 6000)}


def _trap_swarm(rng: random.Random, i: int) -> Dict[str, Any]:
    if i % 2 == 0:  # fake seeders that never complete
        return {"tracker_known": True, "tracker_seeders": rng.randint(50, 600), "tracker_leechers": rng.randint(0, 30),
                "tracker_completed": 0}
    return {"tracker_known": True, "tracker_seeders": 0, "tracker_leechers": rng.randint(12, 300),
            "tracker_completed": rng.randint(0, 2), "prev_seeders": rng.randint(10, 80)}


def _media_pool(seed: int, n: int) -> List[Dict[str, Any]]:
    """Verified library imports: bare-title miniseries packs, BluRay movies, and SxxEyy episodes."""
    rng = random.Random(seed)
    out = []
    for i in range(n):
        title = f"{rng.choice(_A)} {rng.choice(_B)}"
        kind = i % 3
        if kind == 0:
            name = f"{title} (2026) [1080p]"
            files = [{"path": f"{title} (2026)/0{j}.mkv", "size": 2_000_000_000} for j in (1, 2, 3)]
            label = "tv_show"
        elif kind == 1:
            year = rng.randint(1995, 2025)
            name = f"{title.replace(' ', '.')}.{year}.1080p.BluRay.x264-{rng.choice(_G)}"
            files = [{"path": f"{name}.mkv", "size": 9_000_000_000}]
            label = "movie"
        else:
            name = f"{title.replace(' ', '.')}.S0{rng.randint(1, 5)}E0{rng.randint(1, 9)}.1080p.WEB.h264-{rng.choice(_G)}"
            files = [{"path": f"{name}.mkv", "size": 1_500_000_000}]
            label = "tv_show"
        out.append({"title": title, "name": name, "files": files, "label": label,
                    "swarm": _clean_swarm(rng, i), "year": 2026 if kind == 0 else None})
    return out


def _junk_pool(seed: int, n: int) -> List[Dict[str, Any]]:
    """Completed downloads that the physical verifier flags as masquerading executables, with trap swarms."""
    rng = random.Random(seed)
    out = []
    for i in range(n):
        title = f"{rng.choice(_A)}.{rng.choice(_B)}"
        ext = ("mkv.exe", "mp4.scr", "avi.exe")[i % 3]
        name = f"{title}.2026.1080p.WEB-DL.{ext}"
        out.append({"name": name, "files": [{"path": name, "size": 2_400_000}], "swarm": _trap_swarm(rng, i)})
    return out


def _import_media(pipe, agg, item, catalog_id):
    size = int(sum(f["size"] for f in item["files"]))
    row = {"name": item["name"], "size": size, "files": item["files"], "swarm": item["swarm"]}
    if item["label"] == "movie":
        hook = _movie_webhook(item["title"], 2000, catalog_id, item["name"], item["files"])
        return _export(pipe, agg, hook, adapter="movie_webhook", row=row, ffprobe=_ffprobe(size, 6400.0))
    hook = _series_webhook(item["title"], catalog_id, item["name"], item["files"])
    return _export(pipe, agg, hook, adapter="series_webhook", row=row, ffprobe=_ffprobe(size))


def _import_junk(pipe, agg, item):
    row = {"name": item["name"], "size": 2_400_000, "files": item["files"], "swarm": item["swarm"]}
    return _export(pipe, agg, _qbit_completion(item["name"], 2_400_000), adapter="qbittorrent", row=row)


def _adapter_trap_scores(pipe, rows) -> List[float]:
    """The SwarmResidualAdapter's own junk-margin contribution, computed exactly as in `evaluate_swarm`."""
    import torch
    from swarm_priors import extract_way2_source_features, extract_way3_anomaly_features

    logits = pipe._run_neural_logits(rows)
    scores = []
    with torch.no_grad():
        for i, r in enumerate(rows):
            pj = float(pipe.calibrator.calibrate_head("junk", logits["junk"][i].tolist())[1])
            pj = min(max(pj, 1e-6), 1 - 1e-6)
            zl = math.log(pj / (1 - pj))
            out = pipe.swarm_adapter.forward(
                torch.zeros(1, 384), torch.zeros(1, 24), {"junk": torch.tensor([[0.0, zl]])},
                torch.tensor([list(extract_way2_source_features(r, swarm=r["swarm"]))], dtype=torch.float32),
                torch.tensor([list(extract_way3_anomaly_features(r, swarm=r["swarm"]))], dtype=torch.float32),
            )
            scores.append(float(out["junk"][0, 1] - out["junk"][0, 0]) - zl)
    return scores


def _auroc(scores: Sequence[float], labels: Sequence[int]) -> float:
    pos = [s for s, y in zip(scores, labels) if y]
    neg = [s for s, y in zip(scores, labels) if not y]
    return sum((p > n) + 0.5 * (p == n) for p in pos for n in neg) / (len(pos) * len(neg))


def test_2_calibration_and_swarm_adapter_delta_transfer(make_instance, federation):
    instance_b = make_instance()
    agg = RSIUpdateAggregator(epoch=EPOCH)
    peers = [make_instance() for _ in range(6)]  # instance A + 5 peers
    bundles = []
    for i, pipe in enumerate(peers):
        for k, item in enumerate(_media_pool(100 + i, 15)):
            bundles.append(_import_media(pipe, agg, item, 920_000 + 100 * i + k))
        for item in _junk_pool(200 + i, 6):
            bundles.append(_import_junk(pipe, agg, item))
    assert agg.ingest_bundles(bundles) == {"accepted": 126, "rejected": 0}

    cal_delta, cal_stats = agg.fit_calibration_delta(heads=("category",), min_samples=20)
    assert "category" in cal_delta["heads"], cal_stats
    adapter_delta, ad_stats = agg.train_swarm_adapter_delta(instance_b.swarm_adapter, steps=300, seed=0)
    assert adapter_delta is not None, ad_stats
    assert ad_stats["bce_after"] < ad_stats["bce_before"]
    payload = agg.build_update_payload(calibration_delta=cal_delta, swarm_adapter_delta=adapter_delta)

    # Held-out evaluation data that NO instance contributed
    held_media = _media_pool(999, 60)
    held_rows = [{"name": it["name"], "size": sum(f["size"] for f in it["files"]), "files": it["files"]} for it in held_media]
    held_labels = [it["label"] for it in held_media]
    trap_rng = random.Random(4242)
    swarm_eval_rows, swarm_eval_y = [], []
    for j, it in enumerate(_junk_pool(777, 20)):
        swarm_eval_rows.append({"name": it["name"], "size": 2_400_000, "files": it["files"], "swarm": it["swarm"]})
        swarm_eval_y.append(1)
    for j, it in enumerate(_media_pool(778, 20)):
        swarm_eval_rows.append({"name": it["name"], "size": sum(f["size"] for f in it["files"]), "files": it["files"],
                                "swarm": _clean_swarm(trap_rng, j)})
        swarm_eval_y.append(0)
    # Clean-named releases sitting on trap swarms (oracle-labelled; cannot be physically verified)
    clean_named_traps = []
    for j, it in enumerate(_media_pool(779, 20)):
        clean_named_traps.append({"name": it["name"], "size": sum(f["size"] for f in it["files"]), "files": it["files"],
                                  "swarm": _trap_swarm(trap_rng, j)})
    clean_named_y = [1] * 20 + [0] * 20
    clean_named_rows = clean_named_traps + swarm_eval_rows[20:]

    nll_before = instance_b.calibration_nll(held_rows, held_labels, head="category")
    auroc_before = _auroc(_adapter_trap_scores(instance_b, swarm_eval_rows), swarm_eval_y)
    auroc_clean_named_before = _auroc(_adapter_trap_scores(instance_b, clean_named_rows), clean_named_y)
    preds_before = instance_b.classify_batch(swarm_eval_rows + clean_named_traps)
    base_state = {k: v.detach().clone() for k, v in instance_b.swarm_adapter.state_dict().items()}

    attestations, reviews = _attest_all(peers, federation, payload)
    assert [r.status for r in reviews] == ["review_passed"] * 6, [r.reason for r in reviews]
    manifest = RSIUpdateAggregator.sign_manifest(payload, seq=1, signers=federation["signers"][1:],
                                                 version="calib-adapter-2026.10.18", min_asn_quorum=5)
    result = instance_b.stage_and_apply_dht_update(manifest, payload, attestations=attestations, expected_prev_seq=0)
    assert result.applied and result.status == "committed", (result.reason, result.regressions)
    assert sorted(result.components) == ["calibration_delta", "swarm_adapter_safetensors_delta"]

    # Transfer fidelity: B now holds exactly base + aggregator delta, and the fitted temperature
    new_state = instance_b.swarm_adapter.state_dict()
    for name, d in adapter_delta.items():
        expected = base_state[name] + __import__("torch").from_numpy(d)
        assert bool((new_state[name] == expected).all()), name
    untouched = [k for k in base_state if k not in adapter_delta]
    assert all(bool((new_state[k] == base_state[k]).all()) for k in untouched)
    assert instance_b.calibrator.heads["category"]["temperature"] == cal_delta["heads"]["category"]["temperature"]

    nll_after = instance_b.calibration_nll(held_rows, held_labels, head="category")
    auroc_after = _auroc(_adapter_trap_scores(instance_b, swarm_eval_rows), swarm_eval_y)
    auroc_clean_named_after = _auroc(_adapter_trap_scores(instance_b, clean_named_rows), clean_named_y)
    preds_after = instance_b.classify_batch(swarm_eval_rows + clean_named_traps)
    assert nll_after < nll_before
    # Measured result (not tuned): the trained junk-path delta does not reorder held-out traps. The adapter's
    # residual is scaled by the backbone's junk uncertainty (clamped at 0.02 for confident-clean content), so
    # a bounded delta cannot move the ranking. The claim tested here is transfer + non-regression only.
    assert auroc_after >= auroc_before and auroc_clean_named_after >= auroc_clean_named_before

    def _rates(preds):
        verified_style, clean_named = preds[:40], preds[40:]
        return {
            "masquerade_traps_junk": sum(p.junk == "junk" for p in verified_style[:20]),
            "clean_media_junk_false_positives": sum(p.junk == "junk" for p in verified_style[20:]),
            "clean_named_traps_junk": sum(p.junk == "junk" for p in clean_named),
            "clean_named_traps_honeypot": sum(p.honeypot_trap["is_honeypot_trap"] for p in clean_named),
            "categories_changed": None,
        }

    rates_before, rates_after = _rates(preds_before), _rates(preds_after)
    flips = sum(a.category != b.category for a, b in zip(preds_before, preds_after))
    rates_after["categories_changed"] = flips
    assert rates_after["clean_media_junk_false_positives"] <= rates_before["clean_media_junk_false_positives"]

    RESULTS["test_2"] = {
        "contributions": {"bundles": len(bundles), "calibration_profiles": cal_stats["category"]["n"],
                          "swarm_samples": ad_stats["n_samples"], "swarm_junk_samples": ad_stats["n_junk"]},
        "calibration_fit": cal_stats, "adapter_training": ad_stats,
        "held_out_category_nll": {"n": len(held_rows), "before": round(nll_before, 6), "after": round(nll_after, 6)},
        "adapter_trap_auroc_verified_style": {"n": 40, "before": round(auroc_before, 4), "after": round(auroc_after, 4)},
        "adapter_trap_auroc_clean_named": {"n": 40, "before": round(auroc_clean_named_before, 4), "after": round(auroc_clean_named_after, 4)},
        "label_rates_before": rates_before, "label_rates_after": rates_after,
        "payload_sha256": hashlib.sha256(payload).hexdigest(), "payload_bytes": len(payload),
        "receiver_trace": result.verification_trace,
        "receiver_canary_baseline": result.baseline_canary, "receiver_canary_candidate": result.candidate_canary,
        "peer_reviews": [{"status": r.status, "elapsed_ms": r.elapsed_ms,
                          "calibration_check": next((s["detail"] for s in r.verification_trace if s["step"].startswith("ledger_calibration_nll")), None)}
                         for r in reviews],
        "receiver_elapsed_ms": result.elapsed_ms,
    }


# =============================================================================================
# Test 3: poisoned updates are rejected and rolled back; predictions stay bit-identical
# =============================================================================================

def _probe_rows() -> List[Dict[str, Any]]:
    rows = [{"name": _name(t, RECEIVER_VARIANT), "size": 6_900_000_000, "files": _tv_pack_files(t)} for t in TARGETS]
    for it in _junk_pool(31337, 6):
        rows.append({"name": it["name"], "size": 2_400_000, "files": it["files"], "swarm": it["swarm"]})
    for it in _media_pool(31338, 12):
        rows.append({"name": it["name"], "size": sum(f["size"] for f in it["files"]), "files": it["files"], "swarm": it["swarm"]})
    rows.append({"name": "Codec_Update_Required_To_Play.scr", "size": 900_000,
                 "files": [{"path": "Codec_Update_Required_To_Play.scr", "size": 900_000}]})
    rows.append({"name": "Adult_Studio.24.10.18.Scene3.XXX.1080p.MP4-ARC", "size": 1_400_000_000,
                 "files": [{"path": "Adult_Studio.24.10.18.Scene3.XXX.1080p.MP4-ARC.mp4", "size": 1_400_000_000}]})
    return rows


def _poison_adapter_delta(pipe) -> Dict[str, np.ndarray]:
    sd = pipe.swarm_adapter.state_dict()
    enc = np.zeros(tuple(sd["junk_encoder.4.weight"].shape), dtype=np.float32)
    enc[1, :] = 60.0  # drive the junk logit up for every swarm-bearing item
    return {
        "junk_encoder.4.weight": enc,
        "junk_gate.bias": np.array([0.0, 20.0], dtype=np.float32),
        "head_scales.junk": np.array([1.25], dtype=np.float32),
    }


def test_3_poisoned_weight_updates_roll_back_with_identical_predictions(make_instance, federation):
    instance_b = make_instance()
    probe = _probe_rows()
    agg = RSIUpdateAggregator(epoch=EPOCH)
    signers = federation["signers"]

    def snapshot():
        return [p.to_dict() for p in instance_b.classify_batch(probe)], instance_b.dht_state_fingerprint()

    preds0, fp0 = snapshot()
    staged_cases = {
        "calibration_junk_threshold_0.97": agg.build_update_payload(calibration_delta={"mondrian_thresholds": {"junk": 0.97}}),
        "calibration_junk_temperature_6.0": agg.build_update_payload(calibration_delta={"heads": {"junk": {"temperature": 6.0}}}),
        "adapter_junk_flood_delta": agg.build_update_payload(swarm_adapter_delta=_poison_adapter_delta(instance_b)),
    }
    staged_results = {}
    for seq_offset, (case, payload) in enumerate(staged_cases.items(), start=1):
        manifest = RSIUpdateAggregator.sign_manifest(payload, seq=seq_offset, signers=signers[:2],
                                                     version=f"poison-{case}", min_asn_quorum=5)
        res = instance_b.stage_and_apply_dht_update(manifest, payload, attestations=_forged_attestations(federation, payload),
                                                    expected_prev_seq=0)
        # every cryptographic check passed: the colluding quorum is valid, only the canary stands in the way
        assert [s["step"] for s in res.verification_trace if not s["ok"]] == ["golden_canary_non_regression"], res.verification_trace
        assert res.status == "rolled_back" and not res.applied and res.regressions, case
        assert res.seq_after == 0 and res.state_fingerprint_after == fp0
        preds, fp = snapshot()
        assert fp == fp0 and preds == preds0, f"{case}: predictions changed after rollback"
        staged_results[case] = {"reason": res.reason, "regressions": res.regressions,
                                "candidate_canary": res.candidate_canary, "elapsed_ms": res.elapsed_ms}

    # Honest peers refuse to attest a poisoned payload (defence in depth before the receiver's canary)
    honest = make_instance()
    att, review = honest.attest_dht_update(staged_cases["calibration_junk_threshold_0.97"],
                                           signer=federation["peers"][0]["signer"], asn=PEER_ASNS[0])
    assert att is None and review.status == "review_failed"

    # Rejections before staging: nothing is touched
    good = agg.build_update_payload(calibration_delta={"heads": {"category": {"temperature": 0.9}}})
    good_atts = _forged_attestations(federation, good)
    m_good = RSIUpdateAggregator.sign_manifest(good, seq=1, signers=signers[:2], version="v", min_asn_quorum=5)
    sybils = [{"asn": 65000 + i, "signer": generate_federation_signer()} for i in range(6)]
    rogue = generate_federation_signer()
    m_one_sig = RSIUpdateAggregator.sign_manifest(good, seq=1, signers=signers[:1], version="v", min_asn_quorum=5)
    m_rogue_pq = dataclasses.replace(m_good, pq_signatures={rogue[3]: WOTSPlusSHAKE256.sign(rogue[2], m_good.canonical_payload()).hex(),
                                                            signers[1][3]: m_good.pq_signatures[signers[1][3]]})
    canary_override = json.dumps({"schema": "armiger-dht-update/v1", "epoch": EPOCH,
                                  "calibration_delta": {"heads": {"category": {"temperature": 0.9}}},
                                  "golden_canary_override": [{"item_id": "canary_jnk_000", "junk": False}]},
                                 sort_keys=True, separators=(",", ":")).encode()
    pickle_blob = b"\x80\x04\x95" + b"cos\nsystem\n" + b"\x00" * 64
    pickle_st = (len(b'{"w":{"dtype":"F32","shape":[1],"data_offsets":[0,4]}}')).to_bytes(8, "little") + \
        b'{"w":{"dtype":"F32","shape":[1],"data_offsets":[0,4]}}' + b"\x00\x00\x80\x04" + pickle_blob
    nan_adapter = encode_safetensors({"junk_gate.bias": np.array([0.0, np.nan], dtype=np.float32)})
    pre_cases = {
        "payload_tampered_after_signing": (m_good, good.replace(b"0.9", b"0.1"), good_atts, None),
        "only_one_of_three_signers": (m_one_sig, good, good_atts, None),
        "rogue_wots_key_not_pinned": (m_rogue_pq, good, good_atts, None),
        "four_distinct_asns_only": (m_good, good, _forged_attestations(federation, good, federation["peers"][:4]), None),
        "sybil_unregistered_attestors": (m_good, good, _forged_attestations(federation, good, sybils), None),
        "bep44_cas_mismatch": (m_good, good, good_atts, 7),
        "golden_canary_override_smuggled": (RSIUpdateAggregator.sign_manifest(canary_override, seq=1, signers=signers[:2], version="v", min_asn_quorum=5),
                                            canary_override, _forged_attestations(federation, canary_override), None),
        "pickle_bytecode_in_safetensors": (RSIUpdateAggregator.sign_manifest(pickle_st, seq=1, signers=signers[:2], version="v", min_asn_quorum=5, bundle_format="safetensors"),
                                           pickle_st, _forged_attestations(federation, pickle_st), None),
        "nan_adapter_delta": (RSIUpdateAggregator.sign_manifest(nan_adapter, seq=1, signers=signers[:2], version="v", min_asn_quorum=5, bundle_format="safetensors"),
                              nan_adapter, _forged_attestations(federation, nan_adapter), None),
    }
    pre_results = {}
    for case, (manifest, payload, atts, cas) in pre_cases.items():
        res = instance_b.stage_and_apply_dht_update(manifest, payload, attestations=atts, expected_prev_seq=cas)
        assert not res.applied and res.status in ("rejected_verification", "rolled_back"), case
        assert res.seq_after == 0 and instance_b.dht_state_fingerprint() == fp0, case
        pre_results[case] = {"status": res.status, "reason": res.reason}
    assert pre_results["nan_adapter_delta"]["status"] == "rolled_back"  # caught while staging; restored
    assert all(v["status"] == "rejected_verification" for k, v in pre_results.items() if k != "nan_adapter_delta")

    preds_final, fp_final = snapshot()
    assert fp_final == fp0 and preds_final == preds0

    # The same receiver still accepts a legitimate update afterwards (rollback left it healthy)
    ok = instance_b.stage_and_apply_dht_update(m_good, good, attestations=good_atts, expected_prev_seq=0)
    assert ok.applied and ok.seq_after == 1, (ok.reason, ok.regressions)

    RESULTS["test_3"] = {
        "probe_rows": len(probe), "fingerprint_before": fp0, "fingerprint_after_all_rejections": fp_final,
        "predictions_identical": preds_final == preds0,
        "staged_poison_rollbacks": staged_results, "pre_staging_rejections": pre_results,
        "honest_peer_review_of_poison": {"status": review.status, "reason": review.reason},
        "legitimate_update_after_rollbacks": {"status": ok.status, "seq_after": ok.seq_after},
    }