"""Reproduce MTPLX's AR prefix-selection failure without model/GPU work.

Uses the pinned reference's real bank and lookup functions with empty cache
payloads. This verifies selection policy only, not recurrent-state correctness.
"""

import json
from pathlib import Path
from types import SimpleNamespace

from mtplx import generation
from mtplx.session_bank import SessionBank


def main():
    runtime = SimpleNamespace(model_path=Path("/diagnostic/qwen"), mtp_enabled=True)
    bank = SessionBank()
    prompt = list(range(7896))
    for length, policy in ((7460, "cycle"), (7836, "committed")):
        bank.put(runtime=runtime, token_ids=prompt[:length], cache=[],
                 logits=None, hidden=None, mtp_history_policy=policy)

    assert bank.longest_prefix(prompt).prefix_len == 7836
    assert bank.restore(runtime, prompt, mtp_history_policy="cycle",
                        cache_factory=list) is None
    assert bank.last_miss_reason == "policy_mismatch"
    candidates = bank.near_prefix_candidates(prompt, mtp_history_policy="cycle")
    assert [(entry.prefix_len, matched) for entry, matched in candidates] == [
        (7836, 7836), (7460, 7460),
    ]
    assert generation._restore_near_prefix_prompt_state(
        runtime, prompt, base_hidden_variant=None, mtp_hidden_variant=None,
        mtp_history_policy="cycle", session_bank=bank, template_hash=None,
        draft_head_identity=None, policy_fingerprint=None, cache_factory=list,
    ) is None

    # Control: the shorter prefix is usable; merely removing the incompatible
    # longest candidate makes the unchanged exact-restore function serve it.
    del bank._entries[tuple(prompt[:7836])]
    restored = bank.restore(runtime, prompt, mtp_history_policy="cycle",
                            cache_factory=list)
    assert restored is not None and restored.entry.prefix_len == 7460
    print(json.dumps({"ok": True, "reference": "MTPLX SessionBank + generation",
                      "incompatible_longest_prefix": 7836,
                      "shadowed_usable_prefix": 7460, "prompt_tokens": len(prompt),
                      "control_restored_prefix": restored.entry.prefix_len}))


if __name__ == "__main__":
    main()