#!/usr/bin/env python3
"""Stop-token probe: does the checkpoint emit <|im_end|> and terminate?

    MODEL_URL=http://localhost:8000/v1 MODEL=thrasher python3 probe_stop.py

PASS: >=90% finish_reason=="stop" at max_tokens 700 with temp 0.7 sampling
(temp matters: greedy can mask a weak eos row that sampling exposes), zero
leaked control/placeholder tokens in the text.
"""
import json
import os
import sys
import urllib.request

CARD = ("You are Bram Hollis, keeper of the Wayward Lantern inn. Gruff, "
        "observant. Third person, *asterisk action beats*, 1-3 paragraphs. "
        "Stay in character.")

PROMPTS = [
    [{"role": "system", "content": CARD},
     {"role": "user", "content": u}]
    for u in ["*The door bangs open with the storm.* Got room for one more?",
              "What's the story with the lantern this place is named for?",
              "*slides a copper across the bar* Something warm, please.",
              "You hear anything strange from the fen lately?"]
] + [
    [{"role": "user", "content": u}]
    for u in ["Explain the difference between a mutex and a semaphore.",
              "Write a limerick about a lighthouse keeper.",
              "What are three good questions to ask when renting an apartment?",
              "Summarize the plot of Moby-Dick in two sentences."]
]

LEAK_MARKERS = ("<SPECIAL_", "<|im_start|>", "[INST]", "[SYSTEM_PROMPT]", "<s>")


def call(url, model, messages, temp):
    body = json.dumps({"model": model, "messages": messages,
                       "max_tokens": 700, "temperature": temp}).encode()
    req = urllib.request.Request(f"{url}/chat/completions", data=body,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        c = json.loads(r.read())["choices"][0]
    return c["finish_reason"], c["message"]["content"]


def main():
    url = os.environ.get("MODEL_URL", "http://localhost:8000/v1")
    model = os.environ.get("MODEL", "thrasher")
    reps = int(os.environ.get("REPS", "3"))
    stop = length = leaks = 0
    lens = []
    for msgs in PROMPTS:
        for i in range(reps):
            fr, text = call(url, model, msgs, temp=0.7)
            lens.append(len(text))
            if fr == "stop":
                stop += 1
            else:
                length += 1
                print(f"  CEILING ({fr}): {msgs[-1]['content'][:40]!r} -> "
                      f"...{text[-80:]!r}")
            for m in LEAK_MARKERS:
                if m in text:
                    leaks += 1
                    print(f"  LEAK {m!r} in reply to {msgs[-1]['content'][:40]!r}")
    n = stop + length
    rate = stop / n if n else 0.0
    print(f"\nstop-rate: {stop}/{n} = {rate:.0%}   mean len {sum(lens)//len(lens)} chars"
          f"   leaks: {leaks}")
    ok = rate >= 0.9 and leaks == 0
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())