#!/usr/bin/env python3
"""Shared helpers for real Shardok battle benchmark suites."""

from __future__ import annotations

import csv
import json
import os
import sqlite3
import subprocess
import tempfile
import zipfile
from dataclasses import dataclass
from pathlib import Path


REPO = Path(__file__).resolve().parents[8]
BIN_DIR = REPO / "bazel-bin/src/main/cpp/net/eagle0/shardok/ai_battle_simulator"
DEFAULT_EXTRACTOR = BIN_DIR / "real_battle_config_extractor"
DEFAULT_METADATA = BIN_DIR / "real_battle_smoke_metadata"
DEFAULT_SIMULATOR = BIN_DIR / "ai_battle_simulator_main"


@dataclass(frozen=True)
class BattleSource:
    action_seq: int
    battle_blob: bytes
    boundary_action_seq: int
    state_blob: bytes


def run_command(args: list[str], env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(args, text=True, capture_output=True, check=False, env=env)
    if result.returncode != 0:
        raise RuntimeError(
            f"command failed ({result.returncode}): {' '.join(args)}\n"
            f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
        )
    return result


def game_db_from_zip(zip_path: Path) -> tempfile.TemporaryDirectory[str]:
    tmp = tempfile.TemporaryDirectory(prefix="eagle0-real-battle-db-")
    tmp_path = Path(tmp.name)
    with zipfile.ZipFile(zip_path) as archive:
        candidates = [name for name in archive.namelist() if name.endswith("game.db")]
        if not candidates:
            tmp.cleanup()
            raise RuntimeError(f"No game.db found in {zip_path}")
        archive.extract(candidates[0], tmp_path)
        extracted = tmp_path / candidates[0]
        target = tmp_path / "game.db"
        if extracted != target:
            target.write_bytes(extracted.read_bytes())
    return tmp


def read_battle_sources(db: Path) -> list[BattleSource]:
    conn = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    rows = conn.execute(
        """
        SELECT action_seq, new_battle_proto
        FROM action_results
        WHERE new_battle_proto IS NOT NULL
        ORDER BY action_seq
        """
    ).fetchall()
    result = []
    for action_seq, battle_blob in rows:
        state_row = conn.execute(
            """
            SELECT boundary_action_seq, payload
            FROM state_snapshots
            WHERE boundary_action_seq < ?
            ORDER BY boundary_action_seq DESC
            LIMIT 1
            """,
            (action_seq,),
        ).fetchone()
        if state_row is None:
            continue
        result.append(BattleSource(action_seq, battle_blob, state_row[0], state_row[1]))
    conn.close()
    return result


def shardok_results_desc(db: Path, shardok_game_id: str) -> list[bytes]:
    conn = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    rows = conn.execute(
        """
        SELECT payload
        FROM shardok_results
        WHERE shardok_game_id = ?
        ORDER BY action_seq DESC
        """,
        (shardok_game_id,),
    ).fetchall()
    conn.close()
    return [row[0] for row in rows]


def shardok_action_views(db: Path, shardok_game_id: str) -> list[tuple[int, bytes]]:
    conn = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    rows = conn.execute(
        """
        SELECT seq, payload
        FROM shardok_player_results
        WHERE shardok_game_id = ? AND faction_id = -1
        ORDER BY seq
        """,
        (shardok_game_id,),
    ).fetchall()
    conn.close()
    return rows


def side_troop_totals(player_totals: list[dict[str, object]], player_sides: str) -> tuple[int, int]:
    attacker_troops = 0
    defender_troops = 0
    sides = player_sides.split(";")
    for totals in player_totals:
        player_id = int(totals["player_id"])
        troop_count = int(totals["troop_count"])
        if player_id < len(sides) and sides[player_id] == "defender":
            defender_troops += troop_count
        else:
            attacker_troops += troop_count
    return attacker_troops, defender_troops


def parse_sim_csv(stdout: str) -> dict[str, str]:
    lines = [line for line in stdout.splitlines() if line and not line.startswith("ActionPointDistances")]
    header_index = next(i for i, line in enumerate(lines) if line.startswith("config,state_file,"))
    reader = csv.DictReader(lines[header_index : header_index + 2])
    return next(reader)


def write_csv(path: Path, rows: list[dict[str, object]], columns: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore", lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def load_decisive_metadata(
    db: Path,
    metadata_binary: Path,
    battle_path: Path,
    shardok_game_id: str,
    tmpdir: Path,
    action_seq: int,
) -> dict[str, object] | None:
    for index, result_blob in enumerate(shardok_results_desc(db, shardok_game_id)):
        result_path = tmpdir / f"{action_seq}_result_{index}.pb"
        result_path.write_bytes(result_blob)
        metadata = json.loads(
            run_command([str(metadata_binary), f"--battle={battle_path}", f"--result={result_path}"]).stdout
        )
        if metadata["real_winner"] in ("attacker", "defender"):
            return metadata
    return None


def load_real_troop_totals(
    db: Path,
    metadata_binary: Path,
    battle_path: Path,
    metadata: dict[str, object],
    tmpdir: Path,
    action_seq: int,
) -> tuple[object, object]:
    view_paths = []
    for seq, view_blob in shardok_action_views(db, str(metadata["shardok_game_id"])):
        view_path = tmpdir / f"{action_seq}_view_{seq}.pb"
        view_path.write_bytes(view_blob)
        view_paths.append(view_path)
    if not view_paths:
        return "n/a", "n/a"

    view_args = [str(metadata_binary), f"--battle={battle_path}"]
    view_args.extend(f"--action-view={path}" for path in view_paths)
    view_metadata = json.loads(run_command(view_args).stdout)
    return side_troop_totals(view_metadata["player_totals"], str(metadata["player_sides"]))


def extract_config(
    extractor: Path,
    battle_path: Path,
    state_path: Path,
    output_path: Path,
    seed: int,
    max_rounds: int,
    attacker_scoring: str,
    defender_scoring: str,
) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    run_command(
        [
            str(extractor),
            f"--battle={battle_path}",
            f"--state={state_path}",
            f"--output={output_path}",
            f"--seed={seed}",
            f"--max-rounds={max_rounds}",
            f"--attacker-scoring={attacker_scoring}",
            f"--defender-scoring={defender_scoring}",
        ]
    )


def run_simulator(
    simulator: Path,
    config_path: Path,
    trace_path: Path,
    run_label: str,
    experiment_id: int | None,
) -> dict[str, str]:
    trace_path.parent.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    if experiment_id is not None:
        env["SHARDOK_SCORING_EXPERIMENT_ID"] = str(experiment_id)
    else:
        env.pop("SHARDOK_SCORING_EXPERIMENT_ID", None)
    result = run_command(
        [
            str(simulator),
            f"--config={config_path}",
            "--summary=csv",
            "--quiet",
            f"--trace-jsonl={trace_path}",
            f"--run-label={run_label}",
        ],
        env=env,
    )
    return parse_sim_csv(result.stdout)
