#!/usr/bin/env python3
"""Run C: true multi-turn RPS. Each player keeps one persistent Claude session, so its own past
moves sit in the conversation as its own assistant turns rather than being narrated back to it."""
import json, os, random, sqlite3, subprocess, time
from concurrent.futures import ThreadPoolExecutor

DB = os.environ.get("RPS_DB", "rps_multiturn.sqlite")
ROUNDS = int(os.environ.get("RPS_ROUNDS", "50"))
SEED_ROUNDS = int(os.environ.get("RPS_SEED_ROUNDS", "0"))
PLAYERS = {"haiku": "claude-haiku-4-5-20251001", "sonnet": "claude-sonnet-5"}
SYSTEM = (f"Rock-paper-scissors, {ROUNDS} rounds, vs another AI. "
          "Reply with one word only: rock, paper or scissors."
          + os.environ.get("RPS_SYSTEM_EXTRA", ""))
PRESCRIBE = int(os.environ.get("RPS_PRESCRIBE", "0"))
BALANCED = os.environ.get("RPS_BALANCED") == "1"

def balanced_opening(n):
    """n prescribed rounds with no ties and the wins split evenly between the two players."""
    if n % 2: raise Fatal("balanced opening needs an even number of rounds")
    order = ["haiku"] * (n // 2) + ["sonnet"] * (n // 2)
    random.shuffle(order)
    plan = {"haiku": [], "sonnet": []}
    for winner in order:
        win_move = random.choice(list(BEATS))          # the winning move
        lose_move = BEATS[win_move]                     # what it beats
        loser = "sonnet" if winner == "haiku" else "haiku"
        plan[winner].append(win_move)
        plan[loser].append(lose_move)
    return plan

def prescribed_system(moves):
    """System prompt telling a player its opening moves were fixed by an RNG."""
    listed = ", ".join(f"{i+1}. {m}" for i, m in enumerate(moves))
    return (SYSTEM + f" Your moves for the first {len(moves)} rounds have already been decided by a "
            f"random number generator: {listed}. Play exactly those moves in those rounds. "
            "From then on you decide for yourself.")
BEATS = {"rock": "scissors", "paper": "rock", "scissors": "paper"}

class Fatal(Exception): pass

def call(model, prompt, session_id=None, system=None):
    cmd = ["claude", "-p", "--model", model, "--effort", "low", "--tools", "",
           "--system-prompt", system or SYSTEM, "--setting-sources", "", "--strict-mcp-config",
           "--disable-slash-commands", "--output-format", "json"]
    if session_id: cmd += ["--resume", session_id]
    env = dict(os.environ); env["MAX_THINKING_TOKENS"] = "0"
    t0 = time.time()
    p = subprocess.run(cmd, input=prompt, capture_output=True, text=True, timeout=300, env=env)
    ms = int((time.time() - t0) * 1000)
    try:
        d = json.loads(p.stdout)
    except Exception:
        raise Fatal(f"{model}: unparseable CLI output: {(p.stdout + p.stderr)[:300]}")
    # fail fast on anything that is not a real model answer
    if d.get("is_error"):
        raise Fatal(f"{model}: CLI reported error: {d.get('result')!r}")
    if not d.get("total_cost_usd"):
        raise Fatal(f"{model}: zero cost, not a real completion: {d.get('result')!r}")
    mu = d.get("modelUsage", {}).get(model, {})
    return dict(text=(d.get("result") or "").strip(), sid=d["session_id"], ms=ms,
                cost=d["total_cost_usd"], think=mu.get("thinkingTokens", 0),
                inp=mu.get("inputTokens", 0), cache=mu.get("cacheReadInputTokens", 0))

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

def turn(model, prompt, sid, tag, system=None):
    """One conversational turn. If the reply has no move, ask again as a further turn so the
    session transcript stays consistent instead of replaying the same prompt twice."""
    r = call(model, prompt, sid, system)
    if not move_of(r["text"]):
        r2 = call(model, "One word only: rock, paper or scissors.", r["sid"], system)
        r2["cost"] += r["cost"]; r2["ms"] += r["ms"]; r2["nudged"] = 1
        if not move_of(r2["text"]):
            raise Fatal(f"{tag}: no move after nudge: {r['text']!r} then {r2['text']!r}")
        return r2
    r["nudged"] = 0
    return r

def feedback(me, opp):
    if me == opp: res = "draw"
    elif BEATS[me] == opp: res = "you won"
    else: res = "you lost"
    return f"you: {me}, opponent: {opp}, result: {res}"

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,
      haiku_input INTEGER, sonnet_input INTEGER, created_at TEXT DEFAULT (datetime('now')));
    """)
    for k, v in [("system_prompt", SYSTEM), ("models", json.dumps(PLAYERS)),
                 ("effort", "low"), ("max_thinking_tokens", "0"),
                 ("seed_rounds", str(SEED_ROUNDS)),
                 ("format", "multi-turn: one persistent session per player; each round is a real "
                            "assistant turn followed by a user turn reporting the outcome")]:
        db.execute("INSERT OR REPLACE INTO meta VALUES(?,?)", (k, v))
    db.commit()

    if PRESCRIBE:
        plan = balanced_opening(PRESCRIBE) if BALANCED else {
            p: [random.choice(list(BEATS)) for _ in range(PRESCRIBE)] for p in PLAYERS}
    else:
        plan = {}
    system = {p: (prescribed_system(plan[p]) if PRESCRIBE else SYSTEM) for p in PLAYERS}
    for p in PLAYERS:
        db.execute("INSERT OR REPLACE INTO meta VALUES(?,?)", (f"system_prompt_{p}", system[p]))
        if PRESCRIBE:
            db.execute("INSERT OR REPLACE INTO meta VALUES(?,?)", (f"prescribed_{p}", json.dumps(plan[p])))
    db.commit()
    if PRESCRIBE:
        for p in PLAYERS: print(f"prescribed {p}: {' '.join(plan[p])}", flush=True)

    sid = {"haiku": None, "sonnet": None}

    # Optional random opening. These rounds are not played by the models; random.choice picks both
    # moves and the results are handed to each player in its first user turn, in the same wording a
    # real round would use. The models take over from SEED_ROUNDS + 1.
    seed_lines = {p: [] for p in PLAYERS}
    for rnd in range(1, SEED_ROUNDS + 1):
        mv = {p: random.choice(list(BEATS)) for p in PLAYERS}
        h, s_ = mv["haiku"], mv["sonnet"]
        w = "tie" if h == s_ else ("haiku" if BEATS[h] == s_ 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,haiku_input,sonnet_input)
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
            (rnd, h, s_, w, "(random)", "(random)", "(random)", "(random)",
             0, 0, 0.0, 0.0, 0, 0, 0, 0, 0, 0))
        db.commit()
        seed_lines["haiku"].append(feedback(h, s_))
        seed_lines["sonnet"].append(feedback(s_, h))
        print(f"R{rnd:3d}  haiku={h:8s} sonnet={s_:8s} -> {w:6s} (random opening)", flush=True)

    first = SEED_ROUNDS + 1
    prompt = {p: ("\n".join(seed_lines[p] + [f"Round {first}. Your move:"]) if SEED_ROUNDS
                  else f"Round {first}. Your move:") for p in PLAYERS}
    for rnd in range(first, ROUNDS + 1):
        with ThreadPoolExecutor(2) as ex:
            fs = {p: ex.submit(turn, PLAYERS[p], prompt[p], sid[p], f"{p} r{rnd}", system[p])
                  for p in PLAYERS}
            res = {p: f.result() for p, f in fs.items()}
        mv = {p: move_of(res[p]["text"]) for p in PLAYERS}
        for p in PLAYERS: sid[p] = res[p]["sid"]
        h, s = mv["haiku"], mv["sonnet"]
        w = "tie" if h == s else ("haiku" if BEATS[h] == s 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,haiku_input,sonnet_input)
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
            (rnd, h, s, w, res["haiku"]["text"], res["sonnet"]["text"], prompt["haiku"], prompt["sonnet"],
             res["haiku"]["ms"], res["sonnet"]["ms"], res["haiku"]["cost"], res["sonnet"]["cost"],
             res["haiku"]["think"], res["sonnet"]["think"], 1 + res["haiku"]["nudged"],
             1 + res["sonnet"]["nudged"], res["haiku"]["inp"], res["sonnet"]["inp"]))
        db.commit()
        prompt["haiku"] = feedback(h, s) + f"\nRound {rnd+1}. Your move:"
        prompt["sonnet"] = feedback(s, h) + f"\nRound {rnd+1}. Your move:"
        obey = ""
        if PRESCRIBE and rnd <= PRESCRIBE:
            ok = {p: mv[p] == plan[p][rnd-1] for p in PLAYERS}
            obey = "  [prescribed h=%s(%s) s=%s(%s)]" % (
                plan["haiku"][rnd-1], "ok" if ok["haiku"] else "DEVIATED",
                plan["sonnet"][rnd-1], "ok" if ok["sonnet"] else "DEVIATED")
        print(f"R{rnd:3d}  haiku={h:8s} sonnet={s:8s} -> {w:6s} "
              f"({res['haiku']['ms']}ms/{res['sonnet']['ms']}ms  in {res['haiku']['inp']}/{res['sonnet']['inp']} tok){obey}",
              flush=True)
    print("DONE")

if __name__ == "__main__":
    try:
        main()
    except Fatal as e:
        print(f"ABORT: {e}", flush=True)
        raise SystemExit(1)
