"""Probe: is Surjo-50m spec accept 0% a verify bug or genuine drafter behavior?

Compares greedy token streams K=0 vs K in (2,4,8) on the exact Surjo-50m
snapshot, fp32+int8, three prompts (bench / repetitive / paragraph).
Identical streams => verify path is exact. Accept>0 on the repetitive
prompt => drafter+verify path is wired correctly and can accept.

Run from repo root: .venv/Scripts/python scripts/probe_surjo_accept.py
"""

from __future__ import annotations

import sys

from cism.engine import Engine

SURJO = r"G:\hf_cache\models--SurjoLabs--Surjo-50m\snapshots\35c9fa6eacd0e8e977c3e7d878a771837164f930"

PROMPTS: dict[str, str] = {
    "bench": "The future of artificial intelligence is",
    "repetitive": "abc " * 8,
    "paragraph": (
        "The future of artificial intelligence is an interesting subject. "
        "Scientists study how computers learn and how people can use them to solve problems."
    ),
}

PRECISIONS = ("fp32", "int8")
KS = (0, 2, 4, 8)
MAX_NEW_TOKENS = 64
THREADS = 2


def gen(prompt: str, precision: str, spec_k: int):
    eng = Engine.from_pretrained(
        SURJO,
        precision=precision,
        threads=THREADS,
        local_files_only=True,
        trust_remote_code=True,
    )
    return eng.generate(prompt, max_new_tokens=MAX_NEW_TOKENS, temperature=0, spec_k=spec_k)


def main() -> int:
    print(f"model: {SURJO}")
    print(f"precisions={list(PRECISIONS)} K={list(KS)} max_new_tokens={MAX_NEW_TOKENS} threads={THREADS} temp=0")
    for label, text in PROMPTS.items():
        print(f"prompt[{label!r}] = {text!r}")

    rows: list[dict] = []
    failures: list[str] = []
    for pname, ptext in PROMPTS.items():
        for prec in PRECISIONS:
            ref_ids: list[int] | None = None
            ref_text: str | None = None
            for k in KS:
                r = gen(ptext, prec, k)
                if ref_ids is None:
                    ref_ids, ref_text = list(r.token_ids), r.text
                    identical = True
                else:
                    identical = list(r.token_ids) == ref_ids
                    if not identical:
                        failures.append(f"{pname}/{prec}/K={k}: token stream differs from K=0")
                accept_rate = (r.spec_accepted / r.spec_proposed) if r.spec_proposed else None
                rows.append({
                    "prompt": pname, "prec": prec, "K": k,
                    "identical": identical, "proposed": r.spec_proposed,
                    "accepted": r.spec_accepted, "accept_rate": accept_rate,
                    "len": len(r.token_ids), "finish": r.finish_reason,
                })
                rate = f"{accept_rate:6.1%}" if accept_rate is not None else "   n/a"
                print(
                    f"prompt={pname:10s} prec={prec:4s} K={k:2d} "
                    f"identical={str(identical):5s} proposed={r.spec_proposed:3d} "
                    f"accepted={r.spec_accepted:3d} accept={rate} "
                    f"len={len(r.token_ids):3d} finish={r.finish_reason}"
                )
            # show reference text once per (prompt, prec) for behavior context
            preview = (ref_text or "").replace("\n", "\\n")[:160]
            print(f"  ref[{pname}/{prec}] text: {preview!r}")

    print("\nTABLE (prompt, prec, K) -> identical?, proposed, accepted:")
    print("prompt     | prec | K | identical | proposed | accepted | accept% | len | finish")
    for row in rows:
        rate = f"{row['accept_rate']:6.1%}" if row["accept_rate"] is not None else "   n/a"
        print(
            f"{row['prompt']:10s} | {row['prec']:4s} | {row['K']:1d} | "
            f"{str(row['identical']):9s} | {row['proposed']:8d} | {row['accepted']:8d} | "
            f"{rate} | {row['len']:3d} | {row['finish']}"
        )

    any_accept = any(r["accepted"] > 0 for r in rows)
    rep_accept = any(r["prompt"] == "repetitive" and r["accepted"] > 0 for r in rows)
    print(f"\nany cell accepted>0: {any_accept}")
    print(f"repetitive-prompt cell accepted>0: {rep_accept}")

    if failures:
        print("\nVERDICT: BUG — spec changes greedy output:")
        for f in failures:
            print(f"  FAIL {f}")
        return 1
    if rep_accept:
        print("\nVERDICT: GENUINE — K0==Kx exact everywhere; verify path accepts "
              "on repetitive text but the bench/paragraph drafts genuinely miss "
              "(0% is drafter behavior, not a verify bug).")
        return 0
    print("\nVERDICT: INCONCLUSIVE/SUSPICIOUS — K0==Kx exact but nothing ever "
          "accepted, even on repetitive text; verify path may never accept.")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())