"""Pit two exported policies against each other and record the result.

Uses the same in-Godot inference path that ships in the game
(AIShipController + PolicyNetwork), so eval strength = in-game strength.
Episodes are golden-goal: first goal wins, timeout is a draw. Every randomized
starting state is played twice with sides swapped, so physical-team or arena
asymmetry cannot be mistaken for model strength. Results are appended to
eval_history.json — the bot-progress-over-time record.

Example:
    .venv/bin/python evaluate.py ../Game/bots/rookie.json checkpoints/run01/candidate.json --episodes 40
"""

import argparse
import datetime
import json
import os
import pathlib
import subprocess
import tempfile

TRAINING_DIR = pathlib.Path(__file__).resolve().parent
GAME_DIR = TRAINING_DIR.parent / "Game"
TRAINING_SCENE = "res://scenes/training.tscn"
DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"


def run_half(
    godot_bin: str,
    model_a: str,
    model_b: str,
    episodes: int,
    speedup: int,
    seed: int,
    grounded_a: bool = False,
    grounded_b: bool = False,
    team_size: int = 1,
) -> dict:
    cmd = [
        godot_bin,
        "--display-driver",
        "headless",
        "--rendering-method",
        "gl_compatibility",
        "--audio-driver",
        "Dummy",
        "--log-file",
        str(pathlib.Path(tempfile.gettempdir()) / "cosmic-clash-evaluate-godot.log"),
        "--path",
        str(GAME_DIR),
        TRAINING_SCENE,
        "--headless",
        "--disable-render-loop",
        f"--eval_model_a={model_a}",
        f"--eval_model_b={model_b}",
        f"--eval_episodes={episodes}",
        f"--speedup={speedup}",
        f"--env_seed={seed}",
    ]
    if team_size not in (1, 2):
        raise ValueError("team_size must be 1 or 2")
    if team_size == 2:
        cmd.append("--eval_team_size=2")
    # Must match how each model was actually trained (see AIShipController's
    # allow_vertical/allow_pitch_roll) — a grounded pre-generation-4 model
    # never got a reward gradient on these axes, so leaving them unmasked here
    # adds untrained aerial noise its training never had to contend with.
    if grounded_a:
        cmd += ["--eval_allow_vertical_a=false", "--eval_allow_pitch_roll_a=false"]
    if grounded_b:
        cmd += ["--eval_allow_vertical_b=false", "--eval_allow_pitch_roll_b=false"]
    timeout = episodes * 30 / speedup * 3 + 120  # worst case: all draws, plus margin
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    for line in result.stdout.splitlines():
        if line.startswith("EVAL_RESULT "):
            return json.loads(line[len("EVAL_RESULT "):])
    raise RuntimeError(f"No EVAL_RESULT in godot output:\n{result.stdout[-2000:]}\n{result.stderr[-2000:]}")


def evaluate_pair(
    godot_bin: str,
    model_a: str,
    model_b: str,
    episodes: int,
    speedup: int,
    seed: int,
    grounded_a: bool = False,
    grounded_b: bool = False,
    team_size: int = 1,
) -> dict:
    """Replay one seeded state sequence with the models on opposite sides."""
    if episodes < 2 or episodes % 2 != 0:
        raise ValueError("--episodes must be an even number of at least 2 for paired side swaps")
    if team_size not in (1, 2):
        raise ValueError("team_size must be 1 or 2")

    episodes_per_side = episodes // 2
    first = run_half(
        godot_bin, model_a, model_b, episodes_per_side, speedup, seed,
        grounded_a=grounded_a, grounded_b=grounded_b,
        team_size=team_size,
    )
    second = run_half(
        godot_bin, model_b, model_a, episodes_per_side, speedup, seed,
        grounded_a=grounded_b, grounded_b=grounded_a,
        team_size=team_size,
    )

    a_team_0 = {
        "wins_a": first["goals_a"],
        "wins_b": first["goals_b"],
        "draws": first["draws"],
    }
    a_team_1 = {
        "wins_a": second["goals_b"],
        "wins_b": second["goals_a"],
        "draws": second["draws"],
    }
    record = {
        "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
        "model_a": model_a,
        "model_b": model_b,
        "seed": seed,
        "episodes": first["episodes"] + second["episodes"],
        "wins_a": a_team_0["wins_a"] + a_team_1["wins_a"],
        "wins_b": a_team_0["wins_b"] + a_team_1["wins_b"],
        "draws": a_team_0["draws"] + a_team_1["draws"],
        "side_results": {
            "a_team_0": a_team_0,
            "a_team_1": a_team_1,
        },
        "physical_team_wins": {
            "team_0": first["goals_a"] + second["goals_a"],
            "team_1": first["goals_b"] + second["goals_b"],
        },
    }
    record["win_rate_a"] = round(record["wins_a"] / record["episodes"], 3)
    return record


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model_a", help="Path to first exported policy .json")
    parser.add_argument("model_b", help="Path to second exported policy .json")
    parser.add_argument(
        "--episodes", type=int, default=20,
        help="Total episodes; must be even so every seeded state is replayed with sides swapped",
    )
    parser.add_argument(
        "--godot_bin",
        default=os.environ.get("GODOT_BIN", DEFAULT_GODOT_MACOS),
        help="Path to the Godot binary (or set GODOT_BIN)",
    )
    parser.add_argument("--speedup", type=int, default=16)
    parser.add_argument("--team-size", type=int, choices=(1, 2), default=1)
    parser.add_argument("--seed", type=int, default=1, help="Seed for the paired starting-state sequence")
    parser.add_argument("--history", default=str(TRAINING_DIR / "eval_history.json"))
    parser.add_argument(
        "--grounded-a", action="store_true", help="model_a was trained with locomotion masked (pre-generation-4 models)"
    )
    parser.add_argument(
        "--grounded-b", action="store_true", help="model_b was trained with locomotion masked (pre-generation-4 models)"
    )
    args = parser.parse_args()

    model_a = str(pathlib.Path(args.model_a).resolve())
    model_b = str(pathlib.Path(args.model_b).resolve())
    try:
        record = evaluate_pair(
            args.godot_bin, model_a, model_b, args.episodes, args.speedup, args.seed,
            grounded_a=args.grounded_a, grounded_b=args.grounded_b,
            team_size=args.team_size,
        )
    except ValueError as error:
        parser.error(str(error))

    history_path = pathlib.Path(args.history)
    history = json.loads(history_path.read_text()) if history_path.exists() else []
    history.append(record)
    history_path.write_text(json.dumps(history, indent=2) + "\n")

    print(
        f"{pathlib.Path(model_a).name} vs {pathlib.Path(model_b).name} over {record['episodes']} episodes: "
        f"{record['wins_a']}-{record['wins_b']} ({record['draws']} draws), "
        f"win rate A = {record['win_rate_a']:.0%}"
    )
    a_team_0 = record["side_results"]["a_team_0"]
    a_team_1 = record["side_results"]["a_team_1"]
    physical = record["physical_team_wins"]
    print(
        f"Paired seed {record['seed']} side split: "
        f"A as team 0 {a_team_0['wins_a']}-{a_team_0['wins_b']} ({a_team_0['draws']} draws); "
        f"A as team 1 {a_team_1['wins_a']}-{a_team_1['wins_b']} ({a_team_1['draws']} draws); "
        f"physical teams 0-1 = {physical['team_0']}-{physical['team_1']}"
    )
    print(f"Appended to {history_path}")


if __name__ == "__main__":
    main()