#!/usr/bin/env python3
"""Regenerate STANDARD-vs-STANDARD real-battle baseline artifacts."""

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"
HUMAN_FACTIONS = {"3", "4"}
BASELINE_PREFIX = "standard_scoring_all40_post_bugfix_baseline"
REAL_VS_SIM_PREFIX = "standard_scoring_all40_real_vs_sim"
AI_VS_AI_PREFIX = "standard_scoring_ai_vs_ai_smoke"

BASELINE_COLUMNS = [
    "variant",
    "action_seq",
    "config",
    "winner",
    "end_reason",
    "result",
    "total_rounds",
    "total_commands",
    "attacker_ai_commands",
    "defender_ai_commands",
    "attacker_avg_depth",
    "defender_avg_depth",
    "attacker_surviving_units",
    "defender_surviving_units",
    "attacker_surviving_troops",
    "defender_surviving_troops",
    "elapsed_seconds",
]

REAL_VS_SIM_COLUMNS = [
    "action_seq",
    "attacker_factions",
    "attacker_player",
    "defender_factions",
    "defender_player",
    "ai_vs_ai",
    "real_winner",
    "sim_winner",
    "winner_match",
    "real_round",
    "sim_round",
    "real_attacker_troops",
    "real_defender_troops",
    "sim_attacker_troops",
    "sim_defender_troops",
]

AI_VS_AI_COLUMNS = [
    "action_seq",
    "shardok_game_id",
    "attacker_factions",
    "defender_factions",
    "real_winner",
    "sim_winner",
    "winner_match",
    "real_round",
    "sim_round",
    "real_attacker_troops",
    "real_defender_troops",
    "sim_attacker_troops",
    "sim_defender_troops",
]


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 player_label(factions: object) -> str:
    faction_ids = str(factions).replace(",", ";").split(";")
    return "human" if any(faction.strip() in HUMAN_FACTIONS for faction in faction_ids) else "AI"


def winner_counts(rows: list[dict[str, object]], field: str) -> tuple[int, int, int]:
    attacker = sum(1 for row in rows if row[field] == "attacker")
    defender = sum(1 for row in rows if row[field] == "defender")
    unknown = len(rows) - attacker - defender
    return attacker, defender, unknown


def troop_margin(row: dict[str, object]) -> int:
    return int(row["attacker_surviving_troops"]) - int(row["defender_surviving_troops"])


def real_troops_text(row: dict[str, object]) -> str:
    return f"{row['real_attacker_troops']} vs {row['real_defender_troops']}"


def sim_troops_text(row: dict[str, object]) -> str:
    return f"{row['sim_attacker_troops']} vs {row['sim_defender_troops']}"


def write_baseline_md(path: Path, rows: list[dict[str, object]]) -> None:
    attacker_wins, defender_wins, _ = winner_counts(rows, "winner")
    last_player = sum(1 for row in rows if row["end_reason"] == "last_player_standing")
    max_rounds = sum(1 for row in rows if row["end_reason"] == "max_rounds_reached")
    closest = sorted(rows, key=lambda row: abs(troop_margin(row)))[:12]

    lines = [
        "# STANDARD Scoring All-40 Baseline",
        "",
        "This baseline was regenerated after promoting expected-impact archery scoring to STANDARD.",
        "Generated battle configs preserve archery/start-fire capability, per-unit flee capability,",
        "explicit zero-size battalions, no-profession heroes, forced commands, seed `1`, and the",
        "31-round defender timeout.",
        "",
        f"Raw rows are recorded in `{BASELINE_PREFIX}.csv`.",
        "",
        "## Summary",
        "",
        "| Metric | Count |",
        "| --- | ---: |",
        f"| Real battles | {len(rows)} |",
        f"| Valid simulated battles | {len(rows)} |",
        "| Invalid / no recoverable troops | 0 |",
        f"| Attacker wins | {attacker_wins} |",
        f"| Defender wins | {defender_wins} |",
        f"| Last-player-standing endings | {last_player} |",
        f"| Max-round defender endings | {max_rounds} |",
        "",
        "## Closest Battles By Survivor Margin",
        "",
        "Survivor margin is `attacker_surviving_troops - defender_surviving_troops`.",
        "",
        "| Action seq | Winner | End reason | Attacker troops | Defender troops | Margin |",
        "| ---: | --- | --- | ---: | ---: | ---: |",
    ]
    for row in closest:
        margin = troop_margin(row)
        lines.append(
            f"| {row['action_seq']} | {row['winner']} | {row['end_reason']} | "
            f"{row['attacker_surviving_troops']} | {row['defender_surviving_troops']} | {margin:+d} |"
        )
    path.write_text("\n".join(lines) + "\n")


def write_real_vs_sim_md(path: Path, rows: list[dict[str, object]]) -> None:
    known_rows = [row for row in rows if row["real_winner"] in ("attacker", "defender")]
    matches = sum(1 for row in known_rows if row["winner_match"] == "true")
    real_attacker, real_defender, real_unknown = winner_counts(rows, "real_winner")
    sim_attacker, sim_defender, _ = winner_counts(rows, "sim_winner")
    human_battles = sum(
        1 for row in rows if row["attacker_player"] == "human" or row["defender_player"] == "human"
    )

    lines = [
        "# All-40 Real vs STANDARD Simulation",
        "",
        "This table compares the saved-game Shardok winner side with a fresh",
        "STANDARD-vs-STANDARD simulator run from the extracted battle setup.",
        "Faction `3` and faction `4` are marked as human players; all other",
        "factions are marked as AI players.",
        "",
        f"- Battles checked: {len(rows)}",
        f"- Known real outcomes: {len(known_rows)}",
        f"- Winner-side matches on known outcomes: {matches}/{len(known_rows)}",
        f"- Real saved-game winners: {real_attacker} attacker, {real_defender} defender, {real_unknown} unknown",
        f"- Sim baseline winners: {sim_attacker} attacker, {sim_defender} defender",
        f"- Battles involving a human side: {human_battles}",
        "",
        "| Action seq | Attacker | Defender | Real winner | Sim winner | Match | Real troops | Sim troops |",
        "| ---: | --- | --- | --- | --- | --- | ---: | ---: |",
    ]
    for row in rows:
        attacker = f"{row['attacker_factions']} ({row['attacker_player']})"
        defender = f"{row['defender_factions']} ({row['defender_player']})"
        match = "yes" if row["winner_match"] == "true" else "no"
        lines.append(
            f"| {row['action_seq']} | {attacker} | {defender} | {row['real_winner']} | "
            f"{row['sim_winner']} | {match} | {real_troops_text(row)} | {sim_troops_text(row)} |"
        )
    path.write_text("\n".join(lines) + "\n")


def write_ai_vs_ai_md(path: Path, rows: list[dict[str, object]], min_match_rate: float) -> None:
    matched = sum(1 for row in rows if row["winner_match"] == "true")
    rate = matched / len(rows) if rows else 0.0
    lines = [
        "# AI-vs-AI Real Battle Smoke Test",
        "",
        "This smoke test reruns saved battles where neither side used faction `3` or `4`,",
        "then compares STANDARD-vs-STANDARD winner side with the winner side recorded by",
        "the original Shardok battle. It is intentionally a broad similarity check rather",
        "than an exact replay assertion.",
        "",
        f"- Checked battles: {len(rows)}",
        f"- Winner-side matches: {matched}",
        f"- Match rate: {rate:.0%}",
        f"- Required match rate: {min_match_rate:.0%}",
        "",
        "| Action seq | Factions | Real winner | Sim winner | Real troops | Sim troops |",
        "| ---: | --- | --- | --- | ---: | ---: |",
    ]
    for row in rows:
        factions = f"{row['attacker_factions']} vs {row['defender_factions']}"
        lines.append(
            f"| {row['action_seq']} | {factions} | {row['real_winner']} | {row['sim_winner']} | "
            f"{real_troops_text(row)} | {sim_troops_text(row)} |"
        )
    path.write_text("\n".join(lines) + "\n")


def build_rows(
    args: argparse.Namespace,
    db: Path,
    tmpdir: Path,
) -> tuple[list[dict[str, object]], list[dict[str, object]], list[dict[str, object]]]:
    baseline_rows = []
    real_vs_sim_rows = []
    ai_vs_ai_rows = []
    configs_dir = args.work_dir / "configs"
    traces_dir = args.work_dir / "traces"

    for source in suite.read_battle_sources(db):
        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)

        initial_metadata = json.loads(suite.run_command([str(args.metadata), f"--battle={battle_path}"]).stdout)
        metadata = suite.load_decisive_metadata(
            db,
            args.metadata,
            battle_path,
            str(initial_metadata["shardok_game_id"]),
            tmpdir,
            source.action_seq,
        )
        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,
        )
        config_path = configs_dir / f"{source.action_seq}_standard.json"
        trace_path = traces_dir / f"{source.action_seq}_standard.jsonl"
        suite.extract_config(
            args.extractor,
            battle_path,
            state_path,
            config_path,
            args.seed,
            args.max_rounds,
            "STANDARD",
            "STANDARD",
        )
        sim = suite.run_simulator(args.simulator, config_path, trace_path, "standard_baseline", None)
        baseline_rows.append(
            {
                "variant": args.variant,
                "action_seq": source.action_seq,
                "config": config_path,
                "winner": sim["winner"],
                "end_reason": sim["end_reason"],
                "result": sim["result"],
                "total_rounds": sim["total_rounds"],
                "total_commands": sim["total_commands"],
                "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"],
                "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"],
                "elapsed_seconds": sim.get("elapsed_seconds", ""),
            }
        )

        attacker_player = player_label(metadata["attacker_factions"])
        defender_player = player_label(metadata["defender_factions"])
        real_vs_sim_row = {
            "action_seq": source.action_seq,
            "attacker_factions": metadata["attacker_factions"],
            "attacker_player": attacker_player,
            "defender_factions": metadata["defender_factions"],
            "defender_player": defender_player,
            "ai_vs_ai": str(attacker_player == "AI" and defender_player == "AI").lower(),
            "real_winner": metadata["real_winner"],
            "sim_winner": sim["winner"],
            "winner_match": str(metadata["real_winner"] == sim["winner"]).lower(),
            "real_round": metadata["real_round"],
            "sim_round": sim["total_rounds"],
            "real_attacker_troops": real_attacker_troops,
            "real_defender_troops": real_defender_troops,
            "sim_attacker_troops": sim["attacker_surviving_troops"],
            "sim_defender_troops": sim["defender_surviving_troops"],
        }
        real_vs_sim_rows.append(real_vs_sim_row)
        if real_vs_sim_row["ai_vs_ai"] == "true":
            ai_vs_ai_rows.append(
                {
                    "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"],
                    "sim_winner": sim["winner"],
                    "winner_match": real_vs_sim_row["winner_match"],
                    "real_round": metadata["real_round"],
                    "sim_round": sim["total_rounds"],
                    "real_attacker_troops": real_attacker_troops,
                    "real_defender_troops": real_defender_troops,
                    "sim_attacker_troops": sim["attacker_surviving_troops"],
                    "sim_defender_troops": sim["defender_surviving_troops"],
                }
            )
        print(
            f"{source.action_seq}: real={metadata['real_winner']} sim={sim['winner']} "
            f"match={real_vs_sim_row['winner_match']}",
            flush=True,
        )

    return baseline_rows, real_vs_sim_rows, ai_vs_ai_rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--db", type=Path)
    parser.add_argument("--zip", type=Path)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--max-rounds", type=int, default=31)
    parser.add_argument("--variant", default="expected_impact_archery_standard_baseline")
    parser.add_argument("--min-ai-vs-ai-match-rate", type=float, default=0.75)
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--work-dir", type=Path, default=Path("/private/tmp/eagle0-standard-baselines"))
    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)
    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}")

        args.out_dir.mkdir(parents=True, exist_ok=True)
        args.work_dir.mkdir(parents=True, exist_ok=True)
        with tempfile.TemporaryDirectory(prefix="eagle0-standard-baselines-") as tmp:
            baseline_rows, real_vs_sim_rows, ai_vs_ai_rows = build_rows(args, db, Path(tmp))

        if not baseline_rows:
            raise RuntimeError("No saved battles were available for the STANDARD baseline")
        if not ai_vs_ai_rows:
            raise RuntimeError("No AI-vs-AI saved battles were available for the smoke baseline")

        suite.write_csv(args.out_dir / f"{BASELINE_PREFIX}.csv", baseline_rows, BASELINE_COLUMNS)
        suite.write_csv(args.out_dir / f"{REAL_VS_SIM_PREFIX}.csv", real_vs_sim_rows, REAL_VS_SIM_COLUMNS)
        suite.write_csv(args.out_dir / f"{AI_VS_AI_PREFIX}.csv", ai_vs_ai_rows, AI_VS_AI_COLUMNS)
        write_baseline_md(args.out_dir / f"{BASELINE_PREFIX}.md", baseline_rows)
        write_real_vs_sim_md(args.out_dir / f"{REAL_VS_SIM_PREFIX}.md", real_vs_sim_rows)
        write_ai_vs_ai_md(args.out_dir / f"{AI_VS_AI_PREFIX}.md", ai_vs_ai_rows, args.min_ai_vs_ai_match_rate)

        match_rate = sum(1 for row in ai_vs_ai_rows if row["winner_match"] == "true") / len(ai_vs_ai_rows)
        if match_rate < args.min_ai_vs_ai_match_rate:
            raise RuntimeError(
                f"AI-vs-AI smoke match rate {match_rate:.0%} is below {args.min_ai_vs_ai_match_rate:.0%}"
            )
    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)
