#!/usr/bin/env python3
"""Run baseline/candidate Shardok AI experiments on saved real battles."""

from __future__ import annotations

import argparse
import json
import sys
import tempfile
from pathlib import Path

import real_battle_suite as suite


DEFAULT_OUT = Path(__file__).resolve().parent / "experimental_results"


SUMMARY_COLUMNS = [
    "experiment_label",
    "run_label",
    "experiment_id",
    "action_seq",
    "shardok_game_id",
    "attacker_factions",
    "defender_factions",
    "real_winner",
    "real_attacker_troops",
    "real_defender_troops",
    "sim_winner",
    "end_reason",
    "total_rounds",
    "attacker_surviving_units",
    "defender_surviving_units",
    "attacker_surviving_troops",
    "defender_surviving_troops",
    "attacker_ai_commands",
    "defender_ai_commands",
    "attacker_avg_depth",
    "defender_avg_depth",
    "config",
    "trace",
]


COMPARISON_COLUMNS = [
    "experiment_label",
    "experiment_id",
    "action_seq",
    "shardok_game_id",
    "real_winner",
    "baseline_winner",
    "candidate_winner",
    "winner_changed",
    "candidate_matches_real",
    "baseline_attacker_troops",
    "candidate_attacker_troops",
    "attacker_troop_delta",
    "baseline_defender_troops",
    "candidate_defender_troops",
    "defender_troop_delta",
]


def parse_battle_filter(value: str) -> set[int] | None:
    if value == "all":
        return None
    return {int(item.strip()) for item in value.split(",") if item.strip()}


def choose_db(args: argparse.Namespace) -> tuple[Path, tempfile.TemporaryDirectory[str] | None]:
    if args.db is not None:
        return args.db, None
    if args.zip is not None:
        tmp = suite.game_db_from_zip(args.zip)
        return Path(tmp.name) / "game.db", tmp
    raise RuntimeError("Pass --db=/path/to/game.db or --zip=/path/to/save.zip")


def metadata_for_source(
    db: Path,
    source: suite.BattleSource,
    battle_path: Path,
    tmpdir: Path,
    metadata_binary: Path,
) -> dict[str, object] | None:
    metadata = json.loads(suite.run_command([str(metadata_binary), f"--battle={battle_path}"]).stdout)
    decisive = suite.load_decisive_metadata(
        db,
        metadata_binary,
        battle_path,
        str(metadata["shardok_game_id"]),
        tmpdir,
        source.action_seq,
    )
    return decisive


def build_summary_row(
    args: argparse.Namespace,
    run_label: str,
    experiment_id: int | None,
    source: suite.BattleSource,
    metadata: dict[str, object],
    real_attacker_troops: object,
    real_defender_troops: object,
    sim: dict[str, str],
    config_path: Path,
    trace_path: Path,
) -> dict[str, object]:
    return {
        "experiment_label": args.experiment_label,
        "run_label": run_label,
        "experiment_id": "" if experiment_id is None else experiment_id,
        "action_seq": source.action_seq,
        "shardok_game_id": metadata["shardok_game_id"],
        "attacker_factions": metadata["attacker_factions"],
        "defender_factions": metadata["defender_factions"],
        "real_winner": metadata["real_winner"],
        "real_attacker_troops": real_attacker_troops,
        "real_defender_troops": real_defender_troops,
        "sim_winner": sim["winner"],
        "end_reason": sim["end_reason"],
        "total_rounds": sim["total_rounds"],
        "attacker_surviving_units": sim["attacker_surviving_units"],
        "defender_surviving_units": sim["defender_surviving_units"],
        "attacker_surviving_troops": sim["attacker_surviving_troops"],
        "defender_surviving_troops": sim["defender_surviving_troops"],
        "attacker_ai_commands": sim["attacker_ai_commands"],
        "defender_ai_commands": sim["defender_ai_commands"],
        "attacker_avg_depth": sim["attacker_avg_depth"],
        "defender_avg_depth": sim["defender_avg_depth"],
        "config": config_path,
        "trace": trace_path,
    }


def compare_rows(args: argparse.Namespace, baseline: dict[str, object], candidate: dict[str, object]) -> dict[str, object]:
    baseline_attacker = float(baseline["attacker_surviving_troops"])
    candidate_attacker = float(candidate["attacker_surviving_troops"])
    baseline_defender = float(baseline["defender_surviving_troops"])
    candidate_defender = float(candidate["defender_surviving_troops"])
    return {
        "experiment_label": args.experiment_label,
        "experiment_id": args.experiment_id,
        "action_seq": baseline["action_seq"],
        "shardok_game_id": baseline["shardok_game_id"],
        "real_winner": baseline["real_winner"],
        "baseline_winner": baseline["sim_winner"],
        "candidate_winner": candidate["sim_winner"],
        "winner_changed": str(baseline["sim_winner"] != candidate["sim_winner"]).lower(),
        "candidate_matches_real": str(candidate["sim_winner"] == baseline["real_winner"]).lower(),
        "baseline_attacker_troops": baseline_attacker,
        "candidate_attacker_troops": candidate_attacker,
        "attacker_troop_delta": candidate_attacker - baseline_attacker,
        "baseline_defender_troops": baseline_defender,
        "candidate_defender_troops": candidate_defender,
        "defender_troop_delta": candidate_defender - baseline_defender,
    }


def write_report(path: Path, args: argparse.Namespace, summary_rows: list[dict[str, object]], comparison_rows: list[dict[str, object]]) -> None:
    baseline_rows = [row for row in summary_rows if row["run_label"] == "baseline"]
    candidate_rows = [row for row in summary_rows if row["run_label"] == "candidate"]
    baseline_attacker_wins = sum(1 for row in baseline_rows if row["sim_winner"] == "attacker")
    candidate_attacker_wins = sum(1 for row in candidate_rows if row["sim_winner"] == "attacker")
    changed = [row for row in comparison_rows if row["winner_changed"] == "true"]
    candidate_matches = sum(1 for row in comparison_rows if row["candidate_matches_real"] == "true")

    lines = [
        f"# {args.experiment_label}",
        "",
        f"- Experiment id: `{args.experiment_id}`",
        f"- Battles: {len(comparison_rows)}",
        f"- Baseline attacker wins: {baseline_attacker_wins}",
        f"- Candidate attacker wins: {candidate_attacker_wins}",
        f"- Candidate matches real winner: {candidate_matches}/{len(comparison_rows)}",
        f"- Winner changes: {len(changed)}",
        "",
        "## Changed Outcomes",
        "",
        "| Action seq | Real | Baseline | Candidate | Attacker troop delta | Defender troop delta |",
        "| ---: | --- | --- | --- | ---: | ---: |",
    ]
    for row in changed:
        lines.append(
            f"| {row['action_seq']} | {row['real_winner']} | {row['baseline_winner']} | "
            f"{row['candidate_winner']} | {float(row['attacker_troop_delta']):.1f} | "
            f"{float(row['defender_troop_delta']):.1f} |"
        )
    if not changed:
        lines.append("| _none_ | | | | | |")

    lines.extend(
        [
            "",
            "## Battle 6248",
            "",
            "| Real | Baseline | Candidate | Baseline troops | Candidate troops |",
            "| --- | --- | --- | ---: | ---: |",
        ]
    )
    row_6248 = next((row for row in comparison_rows if int(row["action_seq"]) == 6248), None)
    if row_6248 is None:
        lines.append("| _not run_ | | | | |")
    else:
        baseline_troops = f"{float(row_6248['baseline_attacker_troops']):.1f} vs {float(row_6248['baseline_defender_troops']):.1f}"
        candidate_troops = f"{float(row_6248['candidate_attacker_troops']):.1f} vs {float(row_6248['candidate_defender_troops']):.1f}"
        lines.append(
            f"| {row_6248['real_winner']} | {row_6248['baseline_winner']} | "
            f"{row_6248['candidate_winner']} | {baseline_troops} | {candidate_troops} |"
        )

    path.write_text("\n".join(lines) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--db", type=Path)
    parser.add_argument("--zip", type=Path)
    parser.add_argument("--battles", default="all")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--max-rounds", type=int, default=31)
    parser.add_argument("--experiment-id", type=int, required=True)
    parser.add_argument("--experiment-label", required=True)
    parser.add_argument("--baseline-attacker-scoring", default="STANDARD")
    parser.add_argument("--baseline-defender-scoring", default="STANDARD")
    parser.add_argument("--candidate-attacker-scoring", default="EXPERIMENTAL")
    parser.add_argument("--candidate-defender-scoring", default="STANDARD")
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--extractor", type=Path, default=suite.DEFAULT_EXTRACTOR)
    parser.add_argument("--metadata", type=Path, default=suite.DEFAULT_METADATA)
    parser.add_argument("--simulator", type=Path, default=suite.DEFAULT_SIMULATOR)
    parser.add_argument("--jobs", type=int, default=1, help="Reserved; runs sequentially by default")
    args = parser.parse_args()

    db_tmp = None
    try:
        db, db_tmp = choose_db(args)
        if not db.exists():
            raise RuntimeError(f"game.db does not exist: {db}")

        out_dir = args.out_dir / args.experiment_label
        configs_dir = out_dir / "configs"
        traces_dir = out_dir / "traces"
        battle_filter = parse_battle_filter(args.battles)
        summary_rows = []
        comparison_rows = []

        with tempfile.TemporaryDirectory(prefix="eagle0-real-battle-experiment-") as tmp:
            tmpdir = Path(tmp)
            for source in suite.read_battle_sources(db):
                if battle_filter is not None and source.action_seq not in battle_filter:
                    continue
                battle_path = tmpdir / f"{source.action_seq}_battle.pb"
                state_path = tmpdir / f"{source.action_seq}_state_{source.boundary_action_seq}.pb"
                battle_path.write_bytes(source.battle_blob)
                state_path.write_bytes(source.state_blob)

                metadata = metadata_for_source(db, source, battle_path, tmpdir, args.metadata)
                if metadata is None:
                    continue
                real_attacker_troops, real_defender_troops = suite.load_real_troop_totals(
                    db,
                    args.metadata,
                    battle_path,
                    metadata,
                    tmpdir,
                    source.action_seq,
                )

                run_rows = {}
                for run_label, attacker_scoring, defender_scoring, experiment_id in [
                    (
                        "baseline",
                        args.baseline_attacker_scoring,
                        args.baseline_defender_scoring,
                        None,
                    ),
                    (
                        "candidate",
                        args.candidate_attacker_scoring,
                        args.candidate_defender_scoring,
                        args.experiment_id,
                    ),
                ]:
                    config_path = configs_dir / f"{source.action_seq}_{run_label}.json"
                    trace_path = traces_dir / f"{source.action_seq}_{run_label}.jsonl"
                    suite.extract_config(
                        args.extractor,
                        battle_path,
                        state_path,
                        config_path,
                        args.seed,
                        args.max_rounds,
                        attacker_scoring,
                        defender_scoring,
                    )
                    sim = suite.run_simulator(
                        args.simulator,
                        config_path,
                        trace_path,
                        run_label,
                        experiment_id,
                    )
                    row = build_summary_row(
                        args,
                        run_label,
                        experiment_id,
                        source,
                        metadata,
                        real_attacker_troops,
                        real_defender_troops,
                        sim,
                        config_path,
                        trace_path,
                    )
                    summary_rows.append(row)
                    run_rows[run_label] = row

                comparison_rows.append(compare_rows(args, run_rows["baseline"], run_rows["candidate"]))
                print(
                    f"{source.action_seq}: baseline={run_rows['baseline']['sim_winner']} "
                    f"candidate={run_rows['candidate']['sim_winner']}",
                    flush=True,
                )

                if args.limit > 0 and len(comparison_rows) >= args.limit:
                    break

        if not comparison_rows:
            raise RuntimeError("No saved battles matched the requested filters")

        suite.write_csv(out_dir / "summary.csv", summary_rows, SUMMARY_COLUMNS)
        suite.write_csv(out_dir / "comparison.csv", comparison_rows, COMPARISON_COLUMNS)
        write_report(out_dir / "report.md", args, summary_rows, comparison_rows)
    finally:
        if db_tmp is not None:
            db_tmp.cleanup()


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        sys.exit(1)
