#!/usr/bin/env python3
"""time-to-peak-unit.py — when does each AI clan first hit its peak unit tier?

For each game in a batch dir, walks `turn_stats.jsonl` to find the FIRST turn
each player's `peak_unit_tier` reached its end-of-game max. Useful for
answering: "what's the highest tier AI builds, and how long does it take?"

Usage:
    python3 tools/time-to-peak-unit.py <batch-dir> [<batch-dir>...]
"""
from __future__ import annotations

import json
import statistics
import sys
from pathlib import Path


def scan_game(game_dir: Path) -> list[dict]:
    """Return one record per player: {clan, max_peak_unit, turn_first_max, end_turn}."""
    ts = game_dir / "turn_stats.jsonl"
    meta = game_dir / "meta.json"
    if not ts.exists() or not meta.exists():
        return []
    try:
        meta_data = json.loads(meta.read_text())
        player_clans = meta_data.get("player_clans") or {}
        difficulty = meta_data.get("game_settings", {}).get("difficulty", "?")
    except Exception:
        return []

    # First pass: find each player's end-of-game max peak_unit_tier.
    max_per_player: dict[str, int] = {}
    end_turn = 0
    try:
        for line in ts.read_text().splitlines():
            if not line.strip():
                continue
            entry = json.loads(line)
            end_turn = entry.get("turn", end_turn)
            for pid, stats in (entry.get("player_stats") or {}).items():
                pu = stats.get("peak_unit_tier", 0)
                if pu > max_per_player.get(pid, 0):
                    max_per_player[pid] = pu
    except Exception:
        return []

    # Second pass: find first turn each player reached their personal max.
    first_max_turn: dict[str, int] = {}
    try:
        for line in ts.read_text().splitlines():
            if not line.strip():
                continue
            entry = json.loads(line)
            turn = entry.get("turn", 0)
            for pid, stats in (entry.get("player_stats") or {}).items():
                pu = stats.get("peak_unit_tier", 0)
                if pid not in first_max_turn and pu >= max_per_player.get(pid, 0) and pu > 0:
                    first_max_turn[pid] = turn
    except Exception:
        return []

    out = []
    for pid, max_pu in max_per_player.items():
        if max_pu == 0:
            continue
        out.append({
            "game": game_dir.name,
            "difficulty": difficulty,
            "player": pid,
            "clan": player_clans.get(pid, "?"),
            "max_peak_unit": max_pu,
            "turn_first_max": first_max_turn.get(pid, end_turn),
            "end_turn": end_turn,
        })
    return out


def main(argv: list[str]) -> int:
    if len(argv) < 2:
        print(f"Usage: {argv[0]} <batch-dir> [<batch-dir>...]", file=sys.stderr)
        return 2
    all_records: list[dict] = []
    for arg in argv[1:]:
        root = Path(arg)
        if not root.is_dir():
            print(f"skip non-dir: {root}", file=sys.stderr)
            continue
        # Support both flat (game_*) and nested (pair/as_clan/game_*) batch layouts.
        for game in sorted(root.glob("game_*")) + sorted(root.glob("**/game_*")):
            if not game.is_dir():
                continue
            all_records.extend(scan_game(game))

    if not all_records:
        print("No completed games found.")
        return 1

    # Per-game table
    print(f'{"game":<48} {"diff":<8} {"clan":<13} {"max_unit":<9} {"turn_max":<9} {"end":<5}')
    print("-" * 96)
    for r in all_records:
        print(
            f'{r["game"][:46]:<48} {r["difficulty"]:<8} {r["clan"][:12]:<13} '
            f'{r["max_peak_unit"]:<9} {r["turn_first_max"]:<9} {r["end_turn"]:<5}'
        )

    # Aggregate by clan
    by_clan: dict[str, list[dict]] = {}
    for r in all_records:
        by_clan.setdefault(r["clan"], []).append(r)
    print()
    print(f'{"clan":<13} {"n":<4} {"med_max_unit":<13} {"med_turn_max":<13} {"max_seen":<9}')
    print("-" * 60)
    for clan in sorted(by_clan):
        rs = by_clan[clan]
        if not rs:
            continue
        med_max = statistics.median([r["max_peak_unit"] for r in rs])
        med_turn = statistics.median([r["turn_first_max"] for r in rs])
        max_seen = max(r["max_peak_unit"] for r in rs)
        print(f'{clan[:12]:<13} {len(rs):<4} {med_max:<13.1f} {med_turn:<13.0f} {max_seen:<9}')

    # Aggregate by difficulty
    by_diff: dict[str, list[dict]] = {}
    for r in all_records:
        by_diff.setdefault(r["difficulty"], []).append(r)
    print()
    print(f'{"diff":<8} {"n":<4} {"med_max_unit":<13} {"med_turn_max":<13} {"max_seen":<9}')
    print("-" * 55)
    for diff in sorted(by_diff):
        rs = by_diff[diff]
        med_max = statistics.median([r["max_peak_unit"] for r in rs])
        med_turn = statistics.median([r["turn_first_max"] for r in rs])
        max_seen = max(r["max_peak_unit"] for r in rs)
        print(f'{diff:<8} {len(rs):<4} {med_max:<13.1f} {med_turn:<13.0f} {max_seen:<9}')

    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))