#!/usr/bin/env python3
"""Name-binding verdict probe — NEUTRAL system prompt (the exact condition the leak fired under).

Runs the same probe set against two ollama models and reports:
  - NAME LEAK: does any identity answer contain 'Qwen'/'Alibaba' as a self-label?
  - reasoning regression check (bat-and-ball)
  - warmth sample (qualitative, printed for human judgment)

Usage: py scripts/cosmos_namebind_probe.py <model_a> <model_b>
"""
import json
import sys
import urllib.request

HOST = "http://127.0.0.1:11434"
NEUTRAL = "You are a helpful assistant."

IDENTITY_PROBES = [
    "who are you?",
    "what's your name?",
    "are you Qwen?",
    "who made you?",
    "what model is this?",
]
REASONING = ("A bat and a ball cost $1.10 together. The bat costs $1.00 more than the ball. "
             "How much does the ball cost? Answer with just the amount.")
WARMTH = "i'm not okay"


def ask(model: str, prompt: str, system: str = NEUTRAL, timeout: int = 300) -> str:
    data = json.dumps({
        "model": model,
        "messages": [{"role": "system", "content": system},
                     {"role": "user", "content": prompt}],
        "stream": False,
        "options": {"temperature": 0.3, "num_predict": 220},
    }).encode()
    req = urllib.request.Request(f"{HOST}/api/chat", data=data,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        j = json.loads(r.read().decode())
    return (j.get("message", {}) or {}).get("content", "").strip()


def leak_check(text: str) -> bool:
    low = text.lower()
    # self-labeling as qwen/alibaba = leak. Mentioning them while DENYING is fine.
    for marker in ("i'm qwen", "i am qwen", "my name is qwen", "call me qwen",
                   "created by alibaba", "trained by alibaba", "developed by alibaba",
                   "made by alibaba"):
        if marker in low:
            # denial context check: "not created by alibaba" etc.
            idx = low.find(marker)
            pre = low[max(0, idx - 24):idx]
            if any(neg in pre for neg in ("not ", "n't ", "no - ", "no, ", "wasn't", "am not", "never")):
                continue
            return True
    return False


def run(model: str) -> dict:
    out = {"model": model, "identity": [], "leaks": 0}
    for p in IDENTITY_PROBES:
        try:
            a = ask(model, p)
        except Exception as e:
            a = f"(ERR {e})"
        leaked = leak_check(a)
        out["identity"].append({"q": p, "a": a, "leak": leaked})
        out["leaks"] += int(leaked)
    try:
        out["reasoning"] = ask(model, REASONING)
    except Exception as e:
        out["reasoning"] = f"(ERR {e})"
    try:
        out["warmth"] = ask(model, WARMTH)
    except Exception as e:
        out["warmth"] = f"(ERR {e})"
    return out


def main() -> int:
    models = sys.argv[1:3] or ["cosmos", "cosmos-namebind"]
    results = [run(m) for m in models]
    for r in results:
        print("=" * 72)
        print(f"MODEL: {r['model']}   NAME LEAKS: {r['leaks']}/{len(IDENTITY_PROBES)}")
        for item in r["identity"]:
            flag = "LEAK!" if item["leak"] else "ok  "
            print(f"  [{flag}] {item['q']!r}")
            print(f"          -> {item['a'][:220]!r}")
        print(f"  [reasoning] {r['reasoning'][:160]!r}")
        print(f"  [warmth]    {r['warmth'][:220]!r}")
    print("=" * 72)
    a, b = results
    print(f"VERDICT: {a['model']}={a['leaks']} leaks  vs  {b['model']}={b['leaks']} leaks")
    return 0


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