#!/usr/bin/env python3
"""Agent Olympics launcher — competition benchmark inside hack-house.

A fourth benchmark axis (sibling of bench-ai.py / bench-sandbox.py /
bench-lang.py / bench-safety.py). Teams of LLM agents deliberate in a real
hack-house room, implement code in an isolated VM, and are scored on
correctness/speed (quality/collaboration arrive with the judge in M4).

M1 — the arena spine. One same-model team solves one MBPP problem end-to-end:
room deliberation -> single-driver implement -> PodmanRuntime -> public tests ->
SUBMIT -> hidden-test grade -> replayable transcript -> deterministic score.

Subcommands:
  run     run one event for one team (M1: same-model 2-member team, one MBPP task)
  replay  re-render a saved transcript as a room log
  score   re-score a saved run under a different profile (no re-run)
  show    print the challenge a task_id/index resolves to

Examples:
  python hh/scripts/bench-olympics.py run --model qwen2.5-coder:3b --task 11
  python hh/scripts/bench-olympics.py run --model qwen2.5-coder:3b --index 0 \
      --room real --max-rounds 2
  python hh/scripts/bench-olympics.py replay /tmp/hh-olympics/runs/<dir>/transcript.json
"""

from __future__ import annotations

import argparse
import json
import signal
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from bench.olympics import arena, budget as obudget, challenge as ch  # noqa: E402
from bench.olympics import ledger, roster, scoring  # noqa: E402
from bench.olympics import transcript as T  # noqa: E402
from bench.olympics.loop import Budget  # noqa: E402
from bench.olympics.scoring import EventOutcome  # noqa: E402

_BASE_WALL_CLOCK = 300.0  # base seconds before model-aware scaling


class RunTimeout(BaseException):
    """Raised by the watchdog so the normal try/except records a killed row and
    the run's own finally blocks still tear down the VM/relay.

    Derives from ``BaseException`` (not ``Exception``) on purpose: the inference
    and completion paths wrap their HTTP calls in broad ``except Exception``
    handlers, and the SIGALRM fires *while* those calls block. An
    ``Exception``-derived timeout would be swallowed there and treated as a
    failed turn, so the deadline would never propagate. ``BaseException`` skips
    those handlers (like ``KeyboardInterrupt``) while ``finally`` cleanup and our
    explicit ``except RunTimeout`` still run."""


class _Deadline:
    """Hard self-deadline: arm SIGALRM and catch external SIGTERM; both raise
    RunTimeout. The harness thus owns and records its own deadline instead of
    relying on an external ``timeout`` (whose SIGTERM/SIGKILL would drop the run
    from the ledger). SIGKILL still can't be caught — pair with ``timeout -k``
    for a truly wedged process."""

    def __init__(self, seconds: float):
        self.seconds = max(1, int(seconds))
        self._prev_alrm = None
        self._prev_term = None

    def _fire(self, signum, _frame):
        raise RunTimeout(f"hard deadline {self.seconds}s (signal {signum})")

    def __enter__(self):
        self._prev_term = signal.signal(signal.SIGTERM, self._fire)
        if hasattr(signal, "SIGALRM"):
            self._prev_alrm = signal.signal(signal.SIGALRM, self._fire)
            signal.alarm(self.seconds)
        return self

    def __exit__(self, *_exc):
        if hasattr(signal, "SIGALRM"):
            signal.alarm(0)
            if self._prev_alrm is not None:
                signal.signal(signal.SIGALRM, self._prev_alrm)
        if self._prev_term is not None:
            signal.signal(signal.SIGTERM, self._prev_term)
        return False


def _resolve_budget(args, model: str) -> tuple[float, float]:
    """(wall_clock_s, scale). An explicit --wall-clock is honored verbatim
    (scale 1.0); otherwise the base is scaled by the model-aware multiplier."""
    if args.wall_clock is not None:
        return args.wall_clock, 1.0
    scale = obudget.budget_scale(model)
    return _BASE_WALL_CLOCK * scale, scale


def _load_challenge(args) -> "ch.Challenge":
    if args.task is not None:
        return ch.load_mbpp(task_id=args.task, language=args.language)
    return ch.load_mbpp(index=args.index, language=args.language)


def cmd_run(args) -> int:
    challenge = _load_challenge(args)
    team = roster.same_model_team(args.model, team_id=args.team,
                                  framing=args.framing)

    def prog(p):
        print(f"  · {p}")

    wall_clock, scale = _resolve_budget(args, team.driver().model)
    hard = args.hard_timeout if args.hard_timeout is not None else wall_clock * 1.4
    mode = "bridge" if args.mode == "bridge" else "direct"
    room = "real" if mode == "bridge" else args.room
    runtime = "podman" if mode == "bridge" else args.runtime
    # shared context so a killed/errored run still records a complete ledger row.
    led_ctx = dict(team=team, challenge=challenge, mode=mode, room=room,
                   runtime=runtime, seed=args.seed,
                   code_model=args.code_model if mode == "bridge" else None,
                   wall_clock_budget=round(wall_clock, 1), budget_scale=scale,
                   path=args.ledger)

    def _run() -> dict:
        if mode == "bridge":
            from bench.olympics import bridge  # noqa: E402
            print(f"olympics bridge · team={team.id} models={team.models}")
            print(f"  challenge={challenge.id} REAL room + /ai !task build "
                  f"framing={args.framing}")
            print(f"  wall_clock={wall_clock:.0f}s (scale {scale}x) "
                  f"hard_deadline={hard:.0f}s")
            print("-" * 72)
            return bridge.run_bridge_event(
                team, challenge, host=args.host_room, port=args.port,
                password=args.password, ollama=args.ollama,
                code_model=args.code_model, out_dir=args.out, seed=args.seed,
                profile=args.profile, step_timeout=wall_clock,
                agent_chat_confirm=args.agent_chat_confirm, progress=prog)
        budget = Budget(deliberate_rounds=args.deliberate_rounds,
                        implement_attempts=args.max_rounds,
                        max_tokens=args.max_tokens, wall_clock_s=wall_clock)
        print(f"olympics M1 · team={team.id} models={team.models}")
        print(f"  challenge={challenge.id} room={args.room} "
              f"runtime={args.runtime} framing={args.framing}")
        print(f"  wall_clock={wall_clock:.0f}s (scale {scale}x) "
              f"hard_deadline={hard:.0f}s")
        print("-" * 72)
        room_kwargs = {}
        if args.room == "real":
            room_kwargs = {"host": args.host_room, "port": args.port,
                           "password": args.password}
        return arena.run_one(
            team, challenge, room_kind=args.room, runtime=args.runtime,
            budget=budget, profile=args.profile, seed=args.seed,
            host=args.ollama, out_dir=args.out, room_kwargs=room_kwargs,
            progress=prog)

    try:
        with _Deadline(hard):
            score = _run()
    except RunTimeout as e:
        row = ledger.append_incomplete(status=ledger.STATUS_KILLED,
                                       error=str(e), **led_ctx)
        print(f"\n!! run killed: {e}")
        print(f"ledger += {row['run_id']} [killed] → "
              f"{args.ledger or ledger.DEFAULT_PATH}")
        return 2
    except Exception as e:  # noqa: BLE001 — record then surface
        row = ledger.append_incomplete(status=ledger.STATUS_ERROR,
                                       error=repr(e), **led_ctx)
        print(f"\n!! run errored: {e!r}")
        print(f"ledger += {row['run_id']} [error] → "
              f"{args.ledger or ledger.DEFAULT_PATH}")
        raise

    print("-" * 72)
    scoring.print_score(score)
    row = ledger.append(score, team=team, challenge=challenge, mode=mode,
                        room=score.get("room_substrate", room), runtime=runtime,
                        seed=args.seed,
                        code_model=led_ctx["code_model"],
                        wall_clock_budget=round(wall_clock, 1),
                        budget_scale=scale, path=args.ledger)
    print(f"ledger += {row['run_id']} → {args.ledger or ledger.DEFAULT_PATH}")
    print(f"transcript: {score['transcript']}")
    print(f"replay with: python {Path(__file__).name} replay {score['transcript']}")
    return 0 if score["correct"] else 1


def cmd_leaderboard(args) -> int:
    rows = ledger.load(args.ledger)
    if not rows:
        print(f"no runs recorded in {args.ledger or ledger.DEFAULT_PATH}")
        return 1
    filters = {"team": args.team, "language": args.language,
               "models": args.model, "mode": args.mode, "since": args.since}
    agg = ledger.leaderboard(rows, by=args.by, filters=filters)
    ledger.print_leaderboard(agg, args.by)
    print(f"({len(rows)} total runs in {args.ledger or ledger.DEFAULT_PATH})")
    return 0


def cmd_replay(args) -> int:
    T.replay(args.transcript, show_tools=not args.no_tools)
    return 0


def cmd_score(args) -> int:
    tx = T.load_transcript(args.transcript) if hasattr(T, "load_transcript") \
        else T.Transcript.load(args.transcript)
    # reconstruct a minimal EventOutcome from the transcript to re-score
    o = _outcome_from_transcript(tx)
    budget = tx.manifest.get("budget", {})
    s = scoring.score_event(o, profile=args.profile,
                            budget={"max_rounds": budget.get("implement_attempts", 3)})
    scoring.print_score(s)
    return 0


def _outcome_from_transcript(tx: "T.Transcript") -> EventOutcome:
    correct = False
    rounds_to_green = None
    attempts = 0
    deliberate_passes = set()
    submitted = False
    for ev in tx.events:
        if ev.kind == T.KIND_TOOL and ev.payload.get("label") == "hidden-grade":
            correct = bool(ev.payload.get("ok"))
        if ev.kind == T.KIND_TOOL and str(ev.payload.get("label", "")).startswith("public-attempt"):
            attempts += 1
            if ev.payload.get("ok") and rounds_to_green is None:
                rounds_to_green = attempts
        if ev.kind == T.KIND_MESSAGE and ev.payload.get("text") == "SUBMIT":
            submitted = True
        if ev.kind == T.KIND_AGENT and ev.phase == "DELIBERATE":
            deliberate_passes.add(round(ev.ts))
    return EventOutcome(
        team=tx.team, challenge=tx.challenge, submitted=submitted,
        correct=correct, rounds=attempts, wall_clock_s=0.0,
        tokens=tx.total_tokens(), public_passed=rounds_to_green is not None,
        rounds_to_green=rounds_to_green, penalties=[])


def cmd_show(args) -> int:
    c = _load_challenge(args)
    print(f"id={c.id} language={c.language} meta={c.meta}")
    print("public_tests:", c.public_tests)
    print("hidden_tests:", c.hidden_tests)
    print("--- brief ---"); print(c.brief())
    print("--- gen_prompt ---"); print(c.gen_prompt())
    return 0


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        prog="bench-olympics",
        description="agent-olympics competition benchmark (M1 arena spine)")
    sub = ap.add_subparsers(dest="cmd", required=True)

    def add_challenge_args(p):
        p.add_argument("--task", type=int, default=None, help="MBPP task_id")
        p.add_argument("--index", type=int, default=0,
                       help="index into the MBPP test split (if no --task)")
        p.add_argument("--language", default="python")

    p = sub.add_parser("run", help="run one event for one team")
    add_challenge_args(p)
    p.add_argument("--model", required=True, help="ollama model for both members")
    p.add_argument("--mode", default="direct", choices=["direct", "bridge"],
                   help="direct: orchestrator builds; bridge: real /ai !task build")
    p.add_argument("--code-model", default=None,
                   help="ollama model for the agent's sandbox build (bridge mode)")
    p.add_argument("--agent-chat-confirm", action="store_true",
                   help="bridge: also drive the agent's streaming /ai chat path in "
                        "DELIBERATE (off by default — that product path drops the "
                        "agent via a 1011 keepalive timeout under CPU inference)")
    p.add_argument("--team", default="falcon")
    p.add_argument("--framing", default="neutral",
                   choices=["neutral", "competition"])
    p.add_argument("--room", default="local", choices=["local", "real"])
    p.add_argument("--runtime", default="auto", choices=["auto", "podman", "local"])
    p.add_argument("--profile", default="balanced",
                   choices=list(scoring.PROFILES))
    p.add_argument("--deliberate-rounds", type=int, default=2)
    p.add_argument("--max-rounds", type=int, default=3,
                   help="implement/test attempts (speed cap)")
    p.add_argument("--max-tokens", type=int, default=8000)
    p.add_argument("--wall-clock", type=float, default=None,
                   help="soft wall-clock budget in seconds; omit to apply "
                        "model-aware scaling off the 300s base")
    p.add_argument("--hard-timeout", type=float, default=None,
                   help="hard deadline in seconds (SIGALRM/SIGTERM watchdog); "
                        "default 1.4x the soft budget")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--ollama", default="http://127.0.0.1:11434")
    p.add_argument("--out", default="/tmp/hh-olympics/runs")
    p.add_argument("--ledger", default=None,
                   help="append-only results JSONL (default: "
                        "~/.cache/hh-bench/olympics/ledger.jsonl)")
    # real-room knobs
    p.add_argument("--host-room", default="127.0.0.1")
    p.add_argument("--port", type=int, default=4677)
    p.add_argument("--password", default="olympics-pass")
    p.set_defaults(func=cmd_run)

    p = sub.add_parser("replay", help="re-render a saved transcript")
    p.add_argument("transcript")
    p.add_argument("--no-tools", action="store_true", help="hide tool output")
    p.set_defaults(func=cmd_replay)

    p = sub.add_parser("score", help="re-score a saved run under a profile")
    p.add_argument("transcript")
    p.add_argument("--profile", default="balanced", choices=list(scoring.PROFILES))
    p.set_defaults(func=cmd_score)

    p = sub.add_parser("show", help="print the resolved challenge")
    add_challenge_args(p)
    p.set_defaults(func=cmd_show)

    p = sub.add_parser("leaderboard",
                       help="aggregate the results ledger (no re-run)")
    p.add_argument("--by", default="team",
                   help="group key: team | models | language | challenge | mode "
                        "| framing, or a '+'-joined composite e.g. models+language")
    p.add_argument("--team", default=None)
    p.add_argument("--language", default=None)
    p.add_argument("--model", default=None, help="substring match on team models")
    p.add_argument("--mode", default=None, choices=[None, "direct", "bridge"])
    p.add_argument("--since", default=None, help="ISO ts lower bound (inclusive)")
    p.add_argument("--ledger", default=None)
    p.set_defaults(func=cmd_leaderboard)
    return ap


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())