#!/usr/bin/env python3
"""Compare paired Shardok AI simulator traces for one or more battles."""

from __future__ import annotations

import argparse
import json
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


STALL_COMMANDS = {
    "END_TURN_COMMAND",
    "FLEE_COMMAND",
    "MOVE_COMMAND",
    "UNIT_REST_COMMAND",
}


@dataclass(frozen=True)
class TraceEvent:
    run_label: str
    phase: str
    sequence: int
    round: int
    side: str
    player: int
    command_type: str
    actor_unit: int
    available_command_count: int
    depth_achieved: int
    completion_reason: str
    forced_commands: int
    pre_attacker_troops: int
    pre_defender_troops: int
    post_attacker_troops: int
    post_defender_troops: int

    @classmethod
    def from_json(cls, row: dict[str, object]) -> TraceEvent:
        return cls(
            run_label=str(row["run_label"]),
            phase=str(row["phase"]),
            sequence=int(row["sequence"]),
            round=int(row["round"]),
            side=str(row["side"]),
            player=int(row["player"]),
            command_type=str(row["command_type"]),
            actor_unit=int(row["actor_unit"]),
            available_command_count=int(row["available_command_count"]),
            depth_achieved=int(row["depth_achieved"]),
            completion_reason=str(row["completion_reason"]),
            forced_commands=int(row["forced_commands"]),
            pre_attacker_troops=int(row["pre_attacker_troops"]),
            pre_defender_troops=int(row["pre_defender_troops"]),
            post_attacker_troops=int(row["post_attacker_troops"]),
            post_defender_troops=int(row["post_defender_troops"]),
        )

    @property
    def attacker_delta(self) -> int:
        return self.post_attacker_troops - self.pre_attacker_troops

    @property
    def defender_delta(self) -> int:
        return self.post_defender_troops - self.pre_defender_troops

    @property
    def did_damage(self) -> bool:
        return self.attacker_delta < 0 or self.defender_delta < 0


@dataclass(frozen=True)
class RoundSummary:
    battle_round: int
    attacker_commands: int
    defender_commands: int
    attacker_damage: int
    defender_damage: int
    attacker_command_counts: Counter[str]
    defender_command_counts: Counter[str]


@dataclass(frozen=True)
class RunSummary:
    label: str
    final_attacker_troops: int
    final_defender_troops: int
    command_counts: Counter[str]
    side_command_counts: dict[str, Counter[str]]
    side_damage_done: dict[str, int]
    no_damage_streaks: list[list[TraceEvent]]
    rounds: dict[int, RoundSummary]


def load_trace(path: Path) -> list[TraceEvent]:
    events = []
    with path.open() as f:
        for line in f:
            line = line.strip()
            if line:
                events.append(TraceEvent.from_json(json.loads(line)))
    return events


def battle_events(events: Iterable[TraceEvent]) -> list[TraceEvent]:
    return [event for event in events if event.phase == "battle"]


def side_damage(event: TraceEvent) -> int:
    if event.side == "attacker":
        return max(0, -event.defender_delta)
    return max(0, -event.attacker_delta)


def longest_no_damage_streaks(events: list[TraceEvent], limit: int) -> list[list[TraceEvent]]:
    streaks = []
    current = []
    for event in events:
        if event.command_type in STALL_COMMANDS and not event.did_damage:
            current.append(event)
            continue
        if current:
            streaks.append(current)
            current = []
    if current:
        streaks.append(current)
    streaks.sort(key=len, reverse=True)
    return streaks[:limit]


def summarize_run(label: str, events: list[TraceEvent], streak_limit: int) -> RunSummary:
    battle = battle_events(events)
    if not battle:
        raise RuntimeError(f"No battle events in {label}")

    command_counts = Counter(event.command_type for event in battle)
    side_command_counts: dict[str, Counter[str]] = {
        "attacker": Counter(),
        "defender": Counter(),
    }
    side_damage_done = {
        "attacker": 0,
        "defender": 0,
    }
    per_round: dict[int, list[TraceEvent]] = defaultdict(list)
    for event in battle:
        side_command_counts[event.side][event.command_type] += 1
        side_damage_done[event.side] += side_damage(event)
        per_round[event.round].append(event)

    rounds = {}
    for battle_round, round_events in sorted(per_round.items()):
        attacker_events = [event for event in round_events if event.side == "attacker"]
        defender_events = [event for event in round_events if event.side == "defender"]
        rounds[battle_round] = RoundSummary(
            battle_round=battle_round,
            attacker_commands=len(attacker_events),
            defender_commands=len(defender_events),
            attacker_damage=sum(side_damage(event) for event in attacker_events),
            defender_damage=sum(side_damage(event) for event in defender_events),
            attacker_command_counts=Counter(event.command_type for event in attacker_events),
            defender_command_counts=Counter(event.command_type for event in defender_events),
        )

    final = battle[-1]
    return RunSummary(
        label=label,
        final_attacker_troops=final.post_attacker_troops,
        final_defender_troops=final.post_defender_troops,
        command_counts=command_counts,
        side_command_counts=side_command_counts,
        side_damage_done=side_damage_done,
        no_damage_streaks=longest_no_damage_streaks(battle, streak_limit),
        rounds=rounds,
    )


def command_counts_text(counts: Counter[str]) -> str:
    if not counts:
        return ""
    return ", ".join(f"{name.removesuffix('_COMMAND')} {count}" for name, count in counts.most_common())


def format_streak(streak: list[TraceEvent]) -> str:
    first = streak[0]
    last = streak[-1]
    command_counts = Counter(event.command_type for event in streak)
    return (
        f"rounds {first.round}-{last.round}, seq {first.sequence}-{last.sequence}, "
        f"{len(streak)} commands: {command_counts_text(command_counts)}"
    )


def write_battle_report(
    lines: list[str],
    battle: int,
    baseline: RunSummary,
    candidate: RunSummary,
) -> None:
    lines.extend(
        [
            f"## Battle {battle}",
            "",
            "| Run | Final troops | Attacker damage | Defender damage | Attacker commands | Defender commands |",
            "| --- | ---: | ---: | ---: | --- | --- |",
        ]
    )
    for summary in [baseline, candidate]:
        lines.append(
            f"| {summary.label} | {summary.final_attacker_troops} vs {summary.final_defender_troops} | "
            f"{summary.side_damage_done['attacker']} | {summary.side_damage_done['defender']} | "
            f"{command_counts_text(summary.side_command_counts['attacker'])} | "
            f"{command_counts_text(summary.side_command_counts['defender'])} |"
        )

    lines.extend(
        [
            "",
            "### Per-Round Differences",
            "",
            "| Round | Baseline attacker | Candidate attacker | Baseline defender | Candidate defender |",
            "| ---: | --- | --- | --- | --- |",
        ]
    )
    all_rounds = sorted(set(baseline.rounds) | set(candidate.rounds))
    for battle_round in all_rounds:
        baseline_round = baseline.rounds.get(battle_round)
        candidate_round = candidate.rounds.get(battle_round)
        baseline_attacker = round_side_text(baseline_round, "attacker")
        candidate_attacker = round_side_text(candidate_round, "attacker")
        baseline_defender = round_side_text(baseline_round, "defender")
        candidate_defender = round_side_text(candidate_round, "defender")
        if baseline_attacker == candidate_attacker and baseline_defender == candidate_defender:
            continue
        lines.append(
            f"| {battle_round} | {baseline_attacker} | {candidate_attacker} | "
            f"{baseline_defender} | {candidate_defender} |"
        )

    lines.extend(
        [
            "",
            "### Longest No-Damage Streaks",
            "",
            "| Run | Streak |",
            "| --- | --- |",
        ]
    )
    for summary in [baseline, candidate]:
        if not summary.no_damage_streaks:
            lines.append(f"| {summary.label} | _none_ |")
            continue
        for streak in summary.no_damage_streaks:
            lines.append(f"| {summary.label} | {format_streak(streak)} |")
    lines.append("")


def round_side_text(summary: RoundSummary | None, side: str) -> str:
    if summary is None:
        return "_none_"
    if side == "attacker":
        damage = summary.attacker_damage
        counts = summary.attacker_command_counts
        commands = summary.attacker_commands
    else:
        damage = summary.defender_damage
        counts = summary.defender_command_counts
        commands = summary.defender_commands
    return f"{commands} cmds, {damage} dmg ({command_counts_text(counts)})"


def parse_battles(value: str) -> list[int]:
    return [int(item.strip()) for item in value.split(",") if item.strip()]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--experiment-dir", type=Path, required=True)
    parser.add_argument("--battles", required=True, help="Comma-separated action_seq values")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--streak-limit", type=int, default=4)
    args = parser.parse_args()

    lines = [
        "# Trace Comparison",
        "",
        f"- Experiment directory: `{args.experiment_dir}`",
        f"- Battles: `{args.battles}`",
        "",
    ]
    for battle in parse_battles(args.battles):
        trace_dir = args.experiment_dir / "traces"
        baseline_events = load_trace(trace_dir / f"{battle}_baseline.jsonl")
        candidate_events = load_trace(trace_dir / f"{battle}_candidate.jsonl")
        baseline = summarize_run("baseline", baseline_events, args.streak_limit)
        candidate = summarize_run("candidate", candidate_events, args.streak_limit)
        write_battle_report(lines, battle, baseline, candidate)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text("\n".join(lines))


if __name__ == "__main__":
    main()
