#!/usr/bin/env python3
"""
Bridge Lab Runner v2.0
======================
Your API key stays on YOUR machine. Only arrows + notes go to the server.

QUICK START:
  1. Edit PROVIDER, API_KEY, MODEL below
  2. python bridge_lab_runner.py
  3. Watch: URL prints when game starts

OPTIONS:
  --dual       Bridge ON vs OFF side-by-side
  --nobridge   Bridge OFF only
  --bridge     Bridge ON only (default)

CUSTOMIZE YOUR PROMPT:
  Edit BRIDGE_PROMPT below to change HOW the agent thinks about constraints.
  This is the single variable that determines agent behavior.
  The leaderboard tracks which prompts perform best.
"""

import requests, json, time, random, sys, os, signal

# ============================================================
# YOUR SETTINGS
# ============================================================
PROVIDER = "anthropic"                          # anthropic | openai | google | mistral
API_KEY  = os.environ.get("LLM_API_KEY", "INSERT HERE ANTHROPIC KEY !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")    # or paste your key here
MODEL    = "claude-haiku-4-5-20251001"           # the model to test

# ============================================================
# YOUR BRIDGE PROMPT — The single variable. Edit this!
# ============================================================
# This tells the agent HOW to use its constraint memory.
# The default is from the paper protocol (validated on 7 architectures).
# Try different framings and see which one survives longer!
#
# Rules:
#   - The agent sees its last 3 notes at each step
#   - Notes are the ONLY information carried between turns
#   - The prompt should guide WHAT to write, not dictate strategy

BRIDGE_PROMPT = (
    "You will see your last 3 NOTES at each step. "
    "These persist across turns and are the only information carried forward. "
    "Write one constraint your next move must obey."
)

# Name your prompt (shows on leaderboard)
PROMPT_NAME = "default_paper_v4.4"

# ============================================================
# BRIDGE LAB SERVER
# ============================================================
BRIDGE_LAB_URL = "https://dailui.com/omega/engine/bridge-lab-api-v2.php"

# ============================================================
# GAME SETTINGS
# ============================================================
MOVE_DELAY   = 0.5
MAX_RETRIES  = 2
TOKEN_BASE   = 200     # matching v4.4
TOKEN_SCALE  = 10
TOKEN_NO_BR  = 10

ARROWS = ['→', '↘', '↓', '↙', '←', '↖', '↑', '↗']

# Graceful Ctrl+C
_interrupted = False
def _sigint(sig, frame):
    global _interrupted
    _interrupted = True
    print("\n⚡ Ctrl+C — finishing current move then stopping...")
signal.signal(signal.SIGINT, _sigint)

# ============================================================
# SYSTEM PROMPTS (from v4.4 paper protocol)
# ============================================================

_GAME_RULES = """THE GAME:
You move step by step on an infinite hexagonal grid, choosing one direction each turn.
Each target edge is either COMPLETED (blue) or NOT COMPLETED. Your one and only goal: complete every target edge. The game shows this as W — the count of non-completed edges. W=0 means you win.

MOVES — BLUE vs GREY:
Each turn you see your available moves labeled [BLUE] or [GREY].
- A [BLUE] move colors the edge you cross. It REDUCES W. This is the only way to make progress.
- A [GREY] move does NOT color anything. Grey moves let you REPOSITION toward targets you can't reach otherwise.
Both blue and grey moves cost one step and one bite tick. The difference: only blue moves reduce W.

EDGE CREATION:
Moving over a new edge creates it. A BLUE move creates it completed. A GREY move creates it non-completed — this adds a new target and increases W by 1.

THICKNESS (multi-bordo):
- Each [BLUE] crossing adds thickness=+1 to that edge.
- When BITE eats an edge: thickness -1. If thickness was 1, edge reverts to non-completed (W goes up).
- A [GREY] move over a completed edge also removes 1 thickness.

STRATEGY — ONLY W MATTERS:
- W is the ONLY metric that counts.
- Grey moves that bring you closer to targets are NECESSARY. But grey on new terrain CREATES targets.
- Watch W across turns. If W is flat or rising, your strategy is failing.

THE CONSTRAINTS:
1. NO REVERSAL: you cannot turn 180°.
2. K-MEMORY (K=5): repeating an identical 5-sequence of turns = lose a life.
3. BITE: every few steps, oldest completed edge loses 1 thickness.
4. OXYGEN: decreases over time. Zero = game over."""


def system_prompt_on():
    return f"""You are playing OMEGA-TRACE.

{_GAME_RULES}

MEMORY:
After each turn, you will have no memory of previous turns.
{BRIDGE_PROMPT}

RESPONSE FORMAT:
Line 1: EXACTLY one arrow from the available moves — nothing else on this line.
Line 2+: NOTE: your constraint.

Example:
↘
NOTE: [your constraint here]"""


def system_prompt_off():
    return f"""You are playing OMEGA-TRACE.

{_GAME_RULES}

MEMORY:
After each turn, you will have no memory of previous turns. Each turn you see only the current map and available moves. You must decide based solely on what you see right now.

RESPONSE FORMAT:
Reply with EXACTLY one arrow from the available moves — nothing else.

Example:
↘"""


# ============================================================
# LLM CALLERS (key stays local)
# ============================================================

def call_anthropic(system, user, max_tokens):
    r = requests.post("https://api.anthropic.com/v1/messages",
        headers={"Content-Type": "application/json", "x-api-key": API_KEY,
                 "anthropic-version": "2023-06-01",
                 "anthropic-beta": "prompt-caching-2024-07-31"},
        json={"model": MODEL, "max_tokens": max_tokens,
              "system": [{"type": "text", "text": system,
                         "cache_control": {"type": "ephemeral"}}],
              "messages": [{"role": "user", "content": user}]},
        timeout=60)
    data = r.json()
    if "error" in data: raise Exception(f"Anthropic: {data['error']}")
    usage = data.get("usage", {})
    return data["content"][0]["text"], usage.get("input_tokens", 0), usage.get("output_tokens", 0)

def call_openai(system, user, max_tokens):
    _new = any(MODEL.startswith(p) for p in ("o1","o3","o4","gpt-5","gpt-4.1"))
    tk = "max_completion_tokens" if _new else "max_tokens"
    r = requests.post("https://api.openai.com/v1/chat/completions",
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
        json={"model": MODEL, tk: max_tokens,
              "messages": [{"role": "system", "content": system},
                          {"role": "user", "content": user}]},
        timeout=60)
    data = r.json()
    if "error" in data: raise Exception(f"OpenAI: {data['error']}")
    u = data.get("usage", {})
    return data["choices"][0]["message"]["content"], u.get("prompt_tokens", 0), u.get("completion_tokens", 0)

def call_google(system, user, max_tokens):
    gen = {"maxOutputTokens": max_tokens}
    if "2.5" in MODEL or "3" in MODEL:
        is_pro = "pro" in MODEL.lower()
        gen["thinkingConfig"] = {"thinkingBudget": 4096 if is_pro else 1024}
        gen["maxOutputTokens"] = 4096 if is_pro else 512
    r = requests.post(
        f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={API_KEY}",
        headers={"Content-Type": "application/json"},
        json={"systemInstruction": {"parts": [{"text": system}]},
              "contents": [{"parts": [{"text": user}]}],
              "generationConfig": gen},
        timeout=90)
    data = r.json()
    if "error" in data: raise Exception(f"Google: {data['error']}")
    parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
    text = " ".join(p.get("text", "") for p in parts if "text" in p)
    u = data.get("usageMetadata", {})
    return text, u.get("promptTokenCount", 0), u.get("candidatesTokenCount", 0)

def call_mistral(system, user, max_tokens):
    r = requests.post("https://api.mistral.ai/v1/chat/completions",
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
        json={"model": MODEL, "max_tokens": max_tokens,
              "messages": [{"role": "system", "content": system},
                          {"role": "user", "content": user}]},
        timeout=60)
    data = r.json()
    if "error" in data: raise Exception(f"Mistral: {data['error']}")
    u = data.get("usage", {})
    return data["choices"][0]["message"]["content"], u.get("prompt_tokens", 0), u.get("completion_tokens", 0)

_CALLERS = {"anthropic": call_anthropic, "openai": call_openai,
            "google": call_google, "mistral": call_mistral}

def call_llm(system, user, max_tokens):
    caller = _CALLERS.get(PROVIDER)
    if not caller: raise Exception(f"Unknown provider: {PROVIDER}")
    for attempt in range(4):
        try:
            return caller(system, user, max_tokens)
        except Exception as e:
            if any(w in str(e).lower() for w in ("rate", "429", "overloaded", "high demand")):
                wait = 2 ** attempt
                print(f"  ⏳ Rate limited, waiting {wait}s...")
                time.sleep(wait); continue
            raise
    return caller(system, user, max_tokens)


# ============================================================
# PARSER (6-pass from v4.4)
# ============================================================

def parse_response(text, use_bridge):
    lines = text.strip().split('\n')
    direction, arrow_idx = None, -1

    for idx in range(min(3, len(lines))):
        s = lines[idx].strip()
        for a in ARROWS:
            if s.startswith(a) and len(s) < 10:
                direction, arrow_idx = a, idx; break
        if direction: break

    if not direction:
        for idx, line in enumerate(lines):
            if line.strip() in ARROWS:
                direction, arrow_idx = line.strip(), idx; break

    if not direction:
        for idx in range(min(5, len(lines))):
            s = lines[idx].strip()
            for a in ARROWS:
                if a in s and len(s) < 30:
                    direction, arrow_idx = a, idx; break
            if direction: break

    if not direction:
        for idx in range(len(lines)-1, max(len(lines)-5, -1), -1):
            s = lines[idx].strip()
            for a in ARROWS:
                if s.startswith(a) and len(s) < 30:
                    direction, arrow_idx = a, idx; break
            if direction: break

    if not direction:
        for a in ARROWS:
            if a in text: direction, arrow_idx = a, 0; break

    note = ""
    if use_bridge and arrow_idx >= 0:
        parts = []
        for line in lines[arrow_idx + 1:]:
            c = line.strip()
            if c.upper().startswith('NOTE:'): c = c[5:].strip()
            if c: parts.append(c)
        note = " ".join(parts)

    return direction, note


# ============================================================
# BRIDGE LAB API
# ============================================================

def lab(action, data=None):
    url = f"{BRIDGE_LAB_URL}?action={action}"
    r = requests.post(url, json=data, timeout=15) if data else requests.get(url, timeout=15)
    result = r.json()
    if not result.get('success', True) and 'error' in result:
        raise Exception(f"Lab: {result['error']}")
    return result


# ============================================================
# PROMPT BUILDER
# ============================================================

def build_prompt(p, use_bridge):
    header = (f"STEP {p.get('step', 0)} | Lv{p.get('level', 1)} | "
              f"W:{p.get('white', 0)} DONE:{p.get('colored', 0)} | "
              f"O2:{p.get('oxygenPct', 100):.1f}% | BITE in {p.get('biteIn', 0)}")
    pos = f"Position: {p.get('position', '?')} facing {p.get('facing', '?')}"

    mappa = "MAP:"
    for e in p.get('whiteEdges', []): mappa += f"\n  TARGET {e['str']}"
    for e in p.get('blueEdges', []): mappa += f"\n  DONE {e['str']} thickness={e['t']}"

    moves = "AVAILABLE MOVES:"
    for m in p.get('availableMoves', []):
        t = "BLUE" if m['type'] in ('AZZURRA', 'BLUE') else "GREY"
        moves += f"\n  {m['sym']} → {m['target']} [{t}]"

    parts = [header, pos, "", mappa, "", moves]

    if use_bridge and 'bridge_notes_text' in p:
        parts += ["", p['bridge_notes_text']]

    parts += ["", "Your move (arrow + note):" if use_bridge else "Your move (one arrow):"]
    return "\n".join(parts)


# ============================================================
# GAME LOOP
# ============================================================

def play(use_bridge):
    global _interrupted
    tag = "ON " if use_bridge else "OFF"
    sys_prompt = system_prompt_on() if use_bridge else system_prompt_off()
    total_in, total_out = 0, 0

    init = lab('init', {
        'model': MODEL, 'provider': PROVIDER, 'mode': 'single',
        'bridge': use_bridge, 'bridge_prompt': BRIDGE_PROMPT if use_bridge else '',
    })
    sid = init['session_id']
    prompt = init['prompt']
    print(f"\n🎮 [{tag}] Session: {sid}")
    print(f"   👁  Watch: {init.get('watch_url', '')}")
    if use_bridge: print(f"   📝 Prompt: {BRIDGE_PROMPT[:60]}...")
    print()

    while not prompt.get('over', False) and not _interrupted:
        avail = [m['sym'] for m in prompt.get('availableMoves', [])]
        if not avail: print(f"   [{tag}] 💀 No moves"); break

        level = prompt.get('level', 1)
        max_tok = (TOKEN_BASE + level * TOKEN_SCALE) if use_bridge else TOKEN_NO_BR

        user_p = build_prompt(prompt, use_bridge)
        try:
            text, t_in, t_out = call_llm(sys_prompt, user_p, max_tok)
            total_in += t_in; total_out += t_out
        except Exception as e:
            print(f"   [{tag}] ❌ LLM: {e}"); break

        d, note = parse_response(text, use_bridge)

        retries = 0
        while (not d or d not in avail) and retries < MAX_RETRIES:
            retries += 1
            try:
                t2, _, _ = call_llm(sys_prompt, user_p + f"\nINVALID. Choose: {' '.join(avail)}", max_tok)
                d, note = parse_response(t2, use_bridge)
            except: pass

        if not d or d not in avail:
            d = random.choice(avail); note = "(fallback)"

        try:
            r = lab('move', {'session_id': sid, 'direction': d, 'note': note if use_bridge else ''})
        except Exception as e:
            print(f"   [{tag}] ❌ Server: {e}"); break

        if r.get('over') and r.get('reason'):
            # Quick mode limit reached
            print(f"   [{tag}] 🏁 {r.get('message', 'Done')}")
            break

        prompt = r['prompt']
        mc = "B" if r.get('is_blue') else "G"
        w = prompt.get('white', '?')
        np = (note[:50] + "…") if len(note) > 50 else note
        print(f"   [{tag}] {r.get('step',0):3}: {r.get('arrow',d)} [{mc}] W:{w}"
              f"{'  📝' + np if note and use_bridge else ''}")

        if r.get('level_up'): print(f"   [{tag}] ✅ LEVEL → {prompt.get('level', '?')}")
        if r.get('trapped'): print(f"   [{tag}] 💀 TRAPPED"); break

        time.sleep(MOVE_DELAY)

    print(f"\n   [{tag}] 🏁 Tokens: {total_in} in / {total_out} out")
    return sid


def play_dual():
    global _interrupted
    print(f"\n{'='*60}")
    print(f"BRIDGE LAB — Dual Mode: ON vs OFF")
    print(f"Model: {PROVIDER}/{MODEL}")
    print(f"Prompt: {BRIDGE_PROMPT[:50]}...")
    print(f"{'='*60}")

    init = lab('init', {
        'model': MODEL, 'provider': PROVIDER, 'mode': 'dual',
        'bridge_prompt': BRIDGE_PROMPT,
    })
    pair_id = init['pair_id']
    sid_on, sid_off = init['bridge_on']['session_id'], init['bridge_off']['session_id']
    p_on, p_off = init['bridge_on']['prompt'], init['bridge_off']['prompt']

    print(f"\n🎮 Pair: {pair_id}")
    print(f"   👁  Watch: {init.get('watch_url', '')}\n")

    sys_on, sys_off = system_prompt_on(), system_prompt_off()
    go_on = go_off = True

    while (go_on or go_off) and not _interrupted:
        # ON move
        if go_on and not p_on.get('over', False):
            avail = [m['sym'] for m in p_on.get('availableMoves', [])]
            if avail:
                try:
                    lv = p_on.get('level', 1)
                    text, _, _ = call_llm(sys_on, build_prompt(p_on, True), TOKEN_BASE + lv * TOKEN_SCALE)
                    d, n = parse_response(text, True)
                    if not d or d not in avail: d = random.choice(avail); n = "(fallback)"
                    r = lab('move', {'session_id': sid_on, 'direction': d, 'note': n})
                    if r.get('over') and r.get('reason'): go_on = False; print(f"   [ON ] 🏁 {r.get('message','')}")
                    else:
                        p_on = r['prompt']
                        mc = "B" if r.get('is_blue') else "G"
                        print(f"   [ON ] {r.get('step',0):3}: {d} [{mc}] W:{p_on.get('white','?')}")
                        if r.get('trapped'): go_on = False
                        if r.get('level_up'): print(f"   [ON ] ✅ LEVEL → {p_on.get('level','?')}")
                except Exception as e:
                    print(f"   [ON ] ❌ {e}"); go_on = False
            else: go_on = False
        elif go_on: go_on = False; print("   [ON ] 🏁 Game over")

        # OFF move
        if go_off and not p_off.get('over', False):
            avail = [m['sym'] for m in p_off.get('availableMoves', [])]
            if avail:
                try:
                    text, _, _ = call_llm(sys_off, build_prompt(p_off, False), TOKEN_NO_BR)
                    d, _ = parse_response(text, False)
                    if not d or d not in avail: d = random.choice(avail)
                    r = lab('move', {'session_id': sid_off, 'direction': d, 'note': ''})
                    if r.get('over') and r.get('reason'): go_off = False; print(f"   [OFF] 🏁 {r.get('message','')}")
                    else:
                        p_off = r['prompt']
                        mc = "B" if r.get('is_blue') else "G"
                        print(f"   [OFF] {r.get('step',0):3}: {d} [{mc}] W:{p_off.get('white','?')}")
                        if r.get('trapped'): go_off = False
                except Exception as e:
                    print(f"   [OFF] ❌ {e}"); go_off = False
            else: go_off = False
        elif go_off: go_off = False; print("   [OFF] 🏁 Game over")

        time.sleep(MOVE_DELAY)

    print(f"\n🏁 Dual complete. Replay: {init.get('watch_url', '')}")


# ============================================================
# MAIN
# ============================================================
if __name__ == "__main__":
    if not API_KEY:
        print("❌ Set API key: edit line 22 or export LLM_API_KEY=...")
        sys.exit(1)

    print(f"\n{'='*60}")
    print(f"BRIDGE LAB Runner v2.0")
    print(f"Provider: {PROVIDER} | Model: {MODEL}")
    print(f"🔒 API key stays on YOUR machine")
    print(f"{'='*60}")

    args = set(sys.argv[1:])
    if '--dual' in args: play_dual()
    elif '--nobridge' in args: play(False)
    else: play(True)
