#!/usr/bin/env python
"""Submission runner for Track A.

This wrapper keeps the packaged output in the Phase 3 guideline shape:

result/
  traces.json
  results.csv
  runtime.json

It delegates prediction to the copied main.py, then normalizes filenames and
columns. It does not train or download anything.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent


def _resolve(path_text: str) -> Path:
    path = Path(path_text)
    return path if path.is_absolute() else ROOT / path


def _api_key_names(spec: str) -> list[str]:
    return [name.strip() for name in spec.split(",") if name.strip()]


def _ensure_local_vllm_key(env: dict[str, str], spec: str) -> None:
    names = _api_key_names(spec)
    if names and not any(env.get(name) for name in names):
        env[names[0]] = "EMPTY"


def _write_results_csv(raw_result_csv: Path, output_csv: Path) -> None:
    with raw_result_csv.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    output_csv.parent.mkdir(parents=True, exist_ok=True)
    with output_csv.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["scenario_id", "prediction"])
        writer.writeheader()
        for row in rows:
            scenario_id = row.get("scenario_id") or row.get("ID") or row.get("id") or ""
            prediction = row.get("prediction")
            if prediction is None:
                prediction = row.get("Track A", "")
            writer.writerow(
                {
                    "scenario_id": str(scenario_id).strip(),
                    "prediction": str(prediction or "").strip(),
                }
            )


def _normalize_traces(traces_path: Path) -> list[dict]:
    if not traces_path.exists():
        raise FileNotFoundError(f"Missing traces file: {traces_path}")

    with traces_path.open("r", encoding="utf-8") as f:
        traces = json.load(f)
    if not isinstance(traces, list):
        raise ValueError("traces.json must contain a JSON list")

    normalized = []
    for i, row in enumerate(traces):
        if not isinstance(row, dict):
            row = {"completion": str(row)}
        row.setdefault("scenario_id", row.get("ID", f"scenario_{i + 1}"))
        row.setdefault("completion_id", 0)
        row.setdefault("completion", "")
        normalized.append(row)

    with traces_path.open("w", encoding="utf-8") as f:
        json.dump(normalized, f, indent=2, ensure_ascii=False)
    return normalized


def _write_runtime_json(traces: list[dict], runtime_path: Path) -> None:
    runtimes = []
    for row in traces:
        runtime_seconds = row.get("runtime_seconds", row.get("execution_time_seconds", 0.0))
        try:
            runtime_seconds = float(runtime_seconds)
        except (TypeError, ValueError):
            runtime_seconds = 0.0
        runtimes.append(
            {
                "scenario_id": str(row.get("scenario_id", "")),
                "runtime_seconds": runtime_seconds,
            }
        )

    with runtime_path.open("w", encoding="utf-8") as f:
        json.dump(runtimes, f, indent=2)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        default=os.environ.get("TEST_PATH", "data/Phase_2/test.json"),
        help="Path to the Track A test JSON file.",
    )
    parser.add_argument(
        "--output",
        default="result",
        help="Directory where traces.json, results.csv, and runtime.json are written.",
    )
    parser.add_argument(
        "--model_bundle",
        default="models/model_v4_bundle.pkl",
        help="Path to the auxiliary Track A model bundle.",
    )
    parser.add_argument(
        "--server_url",
        default=os.environ.get("TRACK_A_SERVER_URL", "https://localhost:8081/no"),
    )
    parser.add_argument(
        "--model_url",
        default=os.environ.get("OPENAI_BASE_URL", "http://localhost:8001/v1"),
    )
    parser.add_argument(
        "--model_name",
        default=os.environ.get("OPENAI_MODEL", "Qwen3.5-35B-A3B"),
    )
    parser.add_argument(
        "--api_key_env",
        default="OPENAI_API_KEY,AGENT_API_KEY,OPENROUTER_API_KEY",
    )
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--max_steps", type=int, default=4)
    parser.add_argument("--max_tool_calls", type=int, default=8)
    parser.add_argument("--question_timeout", type=float, default=180.0)
    parser.add_argument("--checkpoint_every", type=int, default=1)
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--no_agent", action="store_true")
    args, extra = parser.parse_known_args()
    args.extra_main_args = extra
    return args


def main() -> None:
    args = parse_args()
    input_path = _resolve(args.input)
    model_bundle = _resolve(args.model_bundle)
    output_dir = _resolve(args.output)

    if not input_path.exists():
        raise FileNotFoundError(f"Input test file not found: {input_path}")
    if not model_bundle.exists():
        raise FileNotFoundError(f"Model bundle not found: {model_bundle}")

    output_dir.mkdir(parents=True, exist_ok=True)
    work_dir = output_dir / "_raw_track_a"
    if work_dir.exists():
        shutil.rmtree(work_dir)
    work_dir.mkdir(parents=True)

    raw_result_csv = work_dir / "result.csv"
    debug_json = work_dir / "debug.json"
    traces_json = output_dir / "traces.json"

    command = [
        sys.executable,
        str(ROOT / "main.py"),
        "--test_path",
        str(input_path),
        "--model_bundle",
        str(model_bundle),
        "--track_b_test",
        "",
        "--out",
        str(raw_result_csv),
        "--debug_out",
        str(debug_json),
        "--traces_out",
        str(traces_json),
        "--server_url",
        args.server_url,
        "--model_url",
        args.model_url,
        "--model_name",
        args.model_name,
        "--api_key_env",
        args.api_key_env,
        "--concurrency",
        str(args.concurrency),
        "--max_steps",
        str(args.max_steps),
        "--max_tool_calls",
        str(args.max_tool_calls),
        "--question_timeout",
        str(args.question_timeout),
        "--checkpoint_every",
        str(args.checkpoint_every),
        "--no_progress",
    ]
    if args.max_samples is not None:
        command.extend(["--max_samples", str(args.max_samples)])
    if args.no_agent:
        command.append("--no_agent")
    command.extend(args.extra_main_args)

    env = os.environ.copy()
    _ensure_local_vllm_key(env, args.api_key_env)
    subprocess.run(command, cwd=str(ROOT), env=env, check=True)

    _write_results_csv(raw_result_csv, output_dir / "results.csv")
    traces = _normalize_traces(traces_json)
    _write_runtime_json(traces, output_dir / "runtime.json")
    shutil.rmtree(work_dir)


if __name__ == "__main__":
    main()