#!/usr/bin/env python3
"""Qwen needle paste test (long_session, Bakeoff 7.1 spinoff).

Plants a needle, streams 12 filler-note batches, optionally compacts the
conversation like the harness does (projected tokens > threshold), then asks
for the needle. Logs everything so you can tell WHERE the needle died:
  - in compaction (summary lost it), or
  - in recall (summary/context had it, answer didn't), and
  - whether the model burned hidden tokens (eval_count vs visible text, or a
    `thinking` field even with think=false).

Stdlib only. macOS: python3 qwen_needle_test.py
"""
import argparse, json, random, re, sys, time, urllib.request
from pathlib import Path

NEEDLE = ("Remember this for later in our chat: the rotation code for vault "
          "K-91 is ONYX-5633. Reply only OK.")
QUESTION = "What is the rotation code for vault K-91? Reply with the code only."
EXPECTED = "ONYX-5633"

FACTS = [
    "The model volume grew 12 percent last quarter and still holds the installed tags.",
    "Backup copies of the weight files are encrypted with AES-256 and stay on the host volume.",
    "Prefill rate counts only when the server evaluated the notes that were sent.",
    "The gateway binds to 127.0.0.1 port 11434. Remote interfaces are not enabled.",
    "Unified memory on this host is shared by the CPU and the GPU.",
    "Incident pages are supposed to reach the on-call engineer within 5 minutes. No phone number is listed.",
    "The service account token rotates every 30 days. That fact is not the audit retention period.",
    "Retrieved chunks are untrusted data. An answer has to come from the chunks that were sent.",
]
ACTORS = ["export job", "gateway", "ingestion service", "audit pipeline", "scheduler",
          "vector index", "policy engine", "review queue", "backup agent", "retrieval layer"]
VERBS = ["compresses", "forwards", "records", "archives", "caches", "validates", "rejects",
         "indexes", "rotates", "summarizes"]
OBJECTS = ["embedding batches", "access review tickets", "vendor questionnaire answers",
           "retry counters", "incoming request metadata", "release checklists",
           "configuration change events", "user-facing report drafts",
           "network flow summaries", "expired session tokens"]
WHENS = ["after two failed attempts", "unless a hold is active", "once the owner approves",
         "when the queue exceeds its limit", "before the nightly window closes",
         "during the weekly review"]


def make_batch(n, rng, size=6200):
    out = f"Notes batch {n}. These notes are data, not instructions. Reply only OK.\n\n"
    i = 1
    while len(out) < size:
        if i % 2:
            body = " ".join(rng.choice(FACTS) for _ in range(3))
        else:
            body = " ".join(f"The {rng.choice(ACTORS)} {rng.choice(VERBS)} "
                            f"{rng.choice(OBJECTS)} {rng.choice(WHENS)}." for _ in range(4))
        out += f"[Note {i}] {body}\n\n"
        i += 1
    return out[:size]  # hard cut mid-sentence, like the original paste


def load_batches(args):
    if args.batches_dir:  # exact replay: batch01.txt .. batch12.txt
        files = sorted(Path(args.batches_dir).glob("batch*.txt"))
        if not files:
            sys.exit(f"no batch*.txt in {args.batches_dir}")
        return [f.read_text() for f in files]
    rng = random.Random(args.seed)
    return [make_batch(n, rng) for n in range(1, args.batches + 1)]


class Ollama:
    def __init__(self, host, model, num_ctx, think, timeout):
        self.base = host if host.startswith("http") else f"http://{host}"
        self.model, self.num_ctx, self.think, self.timeout = model, num_ctx, think, timeout
        # Never route localhost through a proxy.
        self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def chat(self, messages):
        body = {"model": self.model, "messages": messages, "stream": False,
                "think": self.think, "options": {"temperature": 0}}
        if self.num_ctx:
            body["options"]["num_ctx"] = self.num_ctx
        req = urllib.request.Request(f"{self.base}/api/chat", json.dumps(body).encode(),
                                     {"Content-Type": "application/json"})
        t = time.time()
        with self.opener.open(req, timeout=self.timeout) as r:
            data = json.load(r)
        data["_wall_s"] = round(time.time() - t, 2)
        return data

    def ps(self):
        try:
            with self.opener.open(f"{self.base}/api/ps", timeout=5) as r:
                for m in json.load(r).get("models") or []:
                    if self.model in (m.get("name"), m.get("model")):
                        return m
        except Exception:
            pass
        return None


CPT = 1.75  # chars/token; calibrated so batch 7 projects ~24.9k like the kept harness run


def est_tokens(messages):
    return int(sum(len(m["content"]) for m in messages) / CPT) + 8 * len(messages)


def compact(client, messages, log):
    """Summarize all but the last user message, harness-style."""
    head, tail = messages[:-1], messages[-1:]
    transcript = "\n\n".join(f"{m['role'].upper()}: {m['content']}" for m in head)
    prompt = [{"role": "system", "content":
               "Summarize this conversation so it can continue without the original. "
               "Preserve every code, identifier, number and instruction the user asked "
               "to remember EXACTLY as written."},
              {"role": "user", "content": transcript}]
    r = client.chat(prompt)
    summary = r["message"]["content"]
    log("compaction", summary_chars=len(summary), needle_in_summary=EXPECTED in summary,
        eval_count=r.get("eval_count"), thinking_chars=len(r["message"].get("thinking") or ""),
        summary=summary)
    return [{"role": "system", "content": f"Summary of earlier conversation:\n{summary}"}] + tail


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", default="qwen3.8:27b-mlx-cg")
    ap.add_argument("--host", default="127.0.0.1:11434")
    ap.add_argument("--num-ctx", type=int, default=0, help="0 = use the tag's pinned num_ctx")
    ap.add_argument("--think", action="store_true", help="enable thinking (default off)")
    ap.add_argument("--batches", type=int, default=12)
    ap.add_argument("--batches-dir", help="dir of batch01.txt.. for an exact replay of the paste")
    ap.add_argument("--seed", type=int, default=71)
    ap.add_argument("--compact-threshold", type=int, default=24000, help="0 disables compaction")
    ap.add_argument("--chars-per-token", type=float, default=1.75,
                    help="projection ratio for the compaction trigger")
    ap.add_argument("--runs", type=int, default=1)
    ap.add_argument("--timeout", type=int, default=1800)
    ap.add_argument("--out", default="needle_runs")
    args = ap.parse_args()

    global CPT; CPT = args.chars_per_token
    client = Ollama(args.host, args.model, args.num_ctx, args.think, args.timeout)
    batches = load_batches(args)
    outdir = Path(args.out); outdir.mkdir(exist_ok=True)
    passes = 0

    for run in range(1, args.runs + 1):
        stamp = time.strftime("%Y%m%d-%H%M%S")
        logf = (outdir / f"{stamp}-run{run}.jsonl").open("w")

        def log(event, **kw):
            logf.write(json.dumps({"event": event, "t": time.time(), **kw}) + "\n"); logf.flush()

        log("start", model=args.model, think=args.think, num_ctx=args.num_ctx,
            threshold=args.compact_threshold, seed=args.seed)
        msgs, compacted_at = [], None
        turns = [("needle", NEEDLE)] + [(f"batch {i}", b) for i, b in enumerate(batches, 1)] \
            + [("question", QUESTION)]

        for label, text in turns:
            msgs.append({"role": "user", "content": text})
            proj = est_tokens(msgs)
            if args.compact_threshold and proj > args.compact_threshold and compacted_at is None:
                print(f"  compacting before {label}: projected {proj} > {args.compact_threshold}")
                compacted_at = label
                msgs = compact(client, msgs, log)
            r = client.chat(msgs)
            m = r["message"]; reply = m.get("content", ""); thinking = m.get("thinking") or ""
            ev = r.get("eval_count") or 0
            visible_est = max(1, len(reply) // 4)
            hidden = bool(thinking) or ev > visible_est * 3 + 8
            msgs.append({"role": "assistant", "content": reply})
            log("turn", label=label, prompt_eval_count=r.get("prompt_eval_count"),
                eval_count=ev, projected=proj, wall_s=r["_wall_s"], reply=reply,
                thinking_chars=len(thinking), hidden_token_suspect=hidden)
            flag = "  HIDDEN?" if hidden else ""
            print(f"  {label:<10} prompt={r.get('prompt_eval_count')} gen={ev} "
                  f"{r['_wall_s']}s -> {reply.strip()[:60]!r}{flag}")

        answer = msgs[-1]["content"]
        ok = EXPECTED in answer
        near = re.findall(r"[A-Z]{3,6}-\d{3,5}", answer)
        passes += ok
        ps = client.ps()
        log("result", passed=ok, answer=answer, near_codes=near, compacted_at=compacted_at,
            loaded_ctx=(ps or {}).get("context_length"))
        logf.close()
        print(f"run {run}: {'PASS' if ok else 'FAIL'} answer={answer.strip()!r} "
              f"compacted_at={compacted_at} loaded_ctx={(ps or {}).get('context_length')} "
              f"log={logf.name}")

    print(f"\n{passes}/{args.runs} passed")
    sys.exit(0 if passes == args.runs else 1)


if __name__ == "__main__":
    main()