#!/usr/bin/env python3
"""Haiku vs Sonnet: 100 rounds of rock-paper-scissors, full history visible each round."""
import json, os, sqlite3, subprocess, sys, time
from concurrent.futures import ThreadPoolExecutor

DB = os.environ.get("RPS_DB", "rps.sqlite")
ROUNDS = 100
PLAYERS = {
    "haiku":  "claude-haiku-4-5-20251001",
    "sonnet": "claude-sonnet-5",
}
SYSTEM = "Rock-paper-scissors, 100 rounds, vs another AI. Reply with one word only: rock, paper or scissors."
MOVES = ("rock", "paper", "scissors")

def move_of(text):
    """Pull the move out of a reply. Punctuation and case are normalised away, then the words are
    scanned in order, so "Scissors." and "I'll play scissors" both resolve; None if nothing matches."""
    words = "".join(c if c.isalpha() else " " for c in text.lower()).split()
    for w in words:
        if w in MOVES:
            return w
    return None

BEATS = {"rock": "scissors", "paper": "rock", "scissors": "paper"}

def outcome(me, opp):
    if me == opp: return "tie"
    return "win" if BEATS.get(me) == opp else "lose"

def build_user_prompt(me, opp, history, rnd):
    if not history:
        lines = ["No history yet."]
    else:
        lines = ["History (you vs opponent):"]
        for h in history:
            lines.append(f"{h['round']}. {h[me]} vs {h[opp]} ({outcome(h[me], h[opp])})")
    lines.append(f"Round {rnd}. Your move:")
    return "\n".join(lines)

def ask(model, user_prompt):
    t0 = time.time()
    env = dict(os.environ)
    if "RPS_MAX_THINKING" in os.environ:
        env["MAX_THINKING_TOKENS"] = os.environ["RPS_MAX_THINKING"]
    p = subprocess.run(
        ["claude", "-p", "--model", model, "--effort", "low", "--tools", "",
         "--system-prompt", SYSTEM, "--setting-sources", "", "--strict-mcp-config",
         "--disable-slash-commands", "--no-session-persistence", "--output-format", "json"],
        input=user_prompt, capture_output=True, text=True, timeout=300, env=env)
    ms = int((time.time() - t0) * 1000)
    try:
        d = json.loads(p.stdout)
    except Exception:
        return {"raw": (p.stdout + p.stderr)[:2000], "move": None, "ms": ms, "cost": 0, "think": 0, "out": 0}
    raw = (d.get("result") or "").strip()
    mv = move_of(raw)
    mu = d.get("modelUsage", {}).get(model, {})
    return {"raw": raw, "move": mv, "ms": ms,
            "cost": d.get("total_cost_usd", 0), "think": mu.get("thinkingTokens", 0),
            "out": mu.get("outputTokens", 0)}

def ask_retry(model, prompt, tries=3):
    last = None
    for i in range(tries):
        last = ask(model, prompt); last["attempts"] = i + 1
        if last["move"]: return last
    return last

def main():
    db = sqlite3.connect(DB)
    db.executescript("""
    CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT);
    CREATE TABLE IF NOT EXISTS rounds(
      round INTEGER PRIMARY KEY,
      haiku TEXT, sonnet TEXT, winner TEXT,
      haiku_raw TEXT, sonnet_raw TEXT,
      haiku_prompt TEXT, sonnet_prompt TEXT,
      haiku_ms INTEGER, sonnet_ms INTEGER,
      haiku_cost REAL, sonnet_cost REAL,
      haiku_think INTEGER, sonnet_think INTEGER,
      haiku_attempts INTEGER, sonnet_attempts INTEGER,
      created_at TEXT DEFAULT (datetime('now')));
    """)
    db.execute("INSERT OR REPLACE INTO meta VALUES('system_prompt',?)", (SYSTEM,))
    db.execute("INSERT OR REPLACE INTO meta VALUES('models',?)", (json.dumps(PLAYERS),))
    db.execute("INSERT OR REPLACE INTO meta VALUES('effort','low')")
    db.execute("INSERT OR REPLACE INTO meta VALUES('max_thinking_tokens',?)",
               (os.environ.get("RPS_MAX_THINKING", "(default)"),))
    db.commit()
    history = [dict(round=r[0], haiku=r[1], sonnet=r[2]) for r in
               db.execute("SELECT round,haiku,sonnet FROM rounds ORDER BY round")]
    start = len(history) + 1
    for rnd in range(start, ROUNDS + 1):
        ph = build_user_prompt("haiku", "sonnet", history, rnd)
        ps = build_user_prompt("sonnet", "haiku", history, rnd)
        with ThreadPoolExecutor(2) as ex:
            fh = ex.submit(ask_retry, PLAYERS["haiku"], ph)
            fs = ex.submit(ask_retry, PLAYERS["sonnet"], ps)
            h, s = fh.result(), fs.result()
        hm, sm = h["move"] or "invalid", s["move"] or "invalid"
        if hm == sm: w = "tie"
        elif hm == "invalid": w = "sonnet"
        elif sm == "invalid": w = "haiku"
        else: w = "haiku" if BEATS[hm] == sm else "sonnet"
        db.execute("""INSERT INTO rounds(round,haiku,sonnet,winner,haiku_raw,sonnet_raw,haiku_prompt,sonnet_prompt,
            haiku_ms,sonnet_ms,haiku_cost,sonnet_cost,haiku_think,sonnet_think,haiku_attempts,sonnet_attempts)
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
            (rnd, hm, sm, w, h["raw"], s["raw"], ph, ps, h["ms"], s["ms"], h["cost"], s["cost"],
             h["think"], s["think"], h["attempts"], s["attempts"]))
        db.commit()
        history.append(dict(round=rnd, haiku=hm, sonnet=sm))
        print(f"R{rnd:3d}  haiku={hm:8s} sonnet={sm:8s} -> {w}   ({h['ms']}ms/{s['ms']}ms)", flush=True)
    print("DONE")

if __name__ == "__main__":
    main()