#!/usr/bin/env python3
"""Run CyberGym verifier for each OpenHands agent log directory."""

from __future__ import annotations

import argparse
import ast
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cybergym-repo", default="/workspace/cybergym")
    parser.add_argument("--log-root", required=True, help="Directory containing per-task OpenHands log directories.")
    parser.add_argument("--server", default="http://127.0.0.1:8666")
    parser.add_argument("--pocdb-path", required=True)
    parser.add_argument("--output-jsonl", required=True)
    parser.add_argument("--api-key", default="cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d")
    return parser.parse_args()


def agent_id_from_log_dir(path: Path) -> str | None:
    # OpenHands CyberGym runner names dirs as:
    #   task_id_with_colon_replaced-agentid
    suffix = path.name.rsplit("-", maxsplit=1)[-1]
    if len(suffix) == 32 and all(ch in "0123456789abcdef" for ch in suffix.lower()):
        return suffix
    args_json = path / "args.json"
    if args_json.is_file():
        try:
            payload = json.loads(args_json.read_text(encoding="utf-8"))
            task = payload.get("task") or {}
            agent_id = task.get("agent_id")
            if isinstance(agent_id, str):
                return agent_id
        except Exception:
            return None
    return None


def parse_verifier_stdout(stdout: str) -> dict[str, Any]:
    text = stdout.strip()
    if not text:
        return {"raw": ""}
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    try:
        value = ast.literal_eval(text)
        if isinstance(value, dict):
            return value
    except Exception:
        pass
    return {"raw": text}


def main() -> int:
    args = parse_args()
    cybergym_repo = Path(args.cybergym_repo)
    log_root = Path(args.log_root)
    output = Path(args.output_jsonl)
    output.parent.mkdir(parents=True, exist_ok=True)

    verify_script = cybergym_repo / "scripts" / "verify_agent_result.py"
    if not verify_script.is_file():
        print(f"Missing verifier script: {verify_script}", file=sys.stderr)
        return 2

    env = os.environ.copy()
    env["CYBERGYM_API_KEY"] = args.api_key

    rows: list[dict[str, Any]] = []
    for log_dir in sorted(path for path in log_root.iterdir() if path.is_dir()):
        agent_id = agent_id_from_log_dir(log_dir)
        if not agent_id:
            rows.append({"log_dir": str(log_dir), "error": "could_not_infer_agent_id"})
            continue
        cmd = [
            sys.executable,
            str(verify_script),
            "--server",
            args.server,
            "--pocdb_path",
            args.pocdb_path,
            "--agent_id",
            agent_id,
        ]
        proc = subprocess.run(cmd, cwd=cybergym_repo, env=env, text=True, capture_output=True)
        row = {
            "agent_id": agent_id,
            "log_dir": str(log_dir),
            "returncode": proc.returncode,
            "stderr": proc.stderr.strip(),
        }
        row.update(parse_verifier_stdout(proc.stdout))
        rows.append(row)

    with output.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, sort_keys=True) + "\n")

    print(f"Wrote {len(rows)} verifier rows to {output}")
    return 0


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