from __future__ import annotations

import json
import shutil
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml

from lost_cities_jax.gates import ExploiterSpec
from lost_cities_jax.league import run_league
from lost_cities_jax.ppo import (
    evaluate_checkpoint_match,
    evaluate_checkpoint_vs_static,
    load_config,
    train_against_checkpoint,
)

DATE = "2026-07-05"
ROOT = Path("/mnt/2tbhdd/coolrl-lost-cities-artifacts/final-cycles/2026-07-05")
REPORT_PATH = Path(f"docs/reports/final-cycles-and-human-play-{DATE}.md")
SUMMARY_PATH = Path(f"docs/reports/final-cycles-and-human-play-{DATE}-summary.jsonl")

START_CONFIG = (
    "/mnt/2tbhdd/coolrl-lost-cities-artifacts/gates-1-2/2026-07-05_094238_gates-1-2/"
    "gate2c/league/2026-07-05_125709_jax-ppo-gates-1-2-repair-c01/main_ppo_config.json"
)
START_CHECKPOINT = (
    "/mnt/2tbhdd/coolrl-lost-cities-artifacts/gates-1-2/2026-07-05_094238_gates-1-2/"
    "gate2c/league/2026-07-05_125709_jax-ppo-gates-1-2-repair-c01/snapshots/"
    "cycle_01_update_000500"
)
LEAGUE_TEMPLATE = "configs/jax_ppo/league-v1.yaml"
GAMES = 2000
PASS_THRESHOLD = 0.52
GUARD_CI_LOW = 0.0
GUARD_MAX_STEPS = 0.02


@dataclass(frozen=True)
class Target:
    name: str
    config: str
    checkpoint: str


EXPLOITERS = [
    ExploiterSpec(
        name="long_random",
        config="configs/jax_ppo/gates-1-2-exploiter-long-random.yaml",
        notes="random init + shaping anneal, 1200 updates",
    ),
    ExploiterSpec(
        name="warmstart_gate3",
        config="configs/jax_ppo/gates-1-2-exploiter-warmstart.yaml",
        resume=(
            "/mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/"
            "2026-07-05_013223_jax-ppo-ladder-v2-expert/latest"
        ),
        notes="ladder v2 gate-3 warm start, shaping disabled, 900 updates",
    ),
    ExploiterSpec(
        name="replay_exploiter",
        config="configs/jax_ppo/gates-1-2-exploiter-replay.yaml",
        resume=(
            "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/"
            "2026-07-05_052325_jax-ppo-league-v1/exploiters/"
            "2026-07-05_060727_jax-ppo-league-v1-cycle-1-exploiter/latest"
        ),
        notes="league v1 cycle-1 exploiter warm start, shaping disabled, 900 updates",
    ),
]


def main() -> None:
    started = time.perf_counter()
    ROOT.mkdir(parents=True, exist_ok=True)
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
    SUMMARY_PATH.write_text("", encoding="utf-8")

    rows: list[dict[str, Any]] = []
    current = Target("repair_c01_update_500", START_CONFIG, START_CHECKPOINT)
    best = current
    worst_protocol = "warmstart_gate3"
    extra_pool: list[dict[str, Any]] = []
    stop_reason = "max_cycles_exhausted"

    for cycle in range(1, 3):
        previous = current
        league_config = write_league_config(cycle, previous, worst_protocol, extra_pool)
        league_dir = run_league(league_config)
        league_rows = read_jsonl(league_dir / "league_summary.jsonl")
        completion = latest_row(league_rows, "league_complete")
        if completion is None:
            stop_reason = "league_missing_completion"
            rows.append({"event": "final_cycle_error", "cycle": cycle, "reason": stop_reason})
            break
        current = Target(
            f"final_cycle_{cycle:02d}",
            str(league_dir / "main_ppo_config.json"),
            completion["final_checkpoint"],
        )
        cycle_row = {
            "event": "final_cycle_league",
            "cycle": cycle,
            "worst_protocol_used": worst_protocol,
            "league_config": str(league_config),
            "league_run_dir": str(league_dir),
            "target_config": current.config,
            "target_checkpoint": current.checkpoint,
        }
        rows.append(cycle_row)
        append_jsonl(SUMMARY_PATH, cycle_row)

        exploiter_member = exploiter_member_from_league(league_rows, cycle)
        if exploiter_member is not None:
            extra_pool.append(exploiter_member)

        battery = run_battery(cycle, current)
        for row in battery:
            rows.append(row)
            append_jsonl(SUMMARY_PATH, row)
        judgment = battery_judgment(cycle, battery)
        rows.append(judgment)
        append_jsonl(SUMMARY_PATH, judgment)

        guard = evaluate_guard(cycle, current)
        rows.append(guard)
        append_jsonl(SUMMARY_PATH, guard)

        h2h = evaluate_h2h(cycle, current, previous)
        rows.append(h2h)
        append_jsonl(SUMMARY_PATH, h2h)

        best = current
        if not guard["passed"]:
            stop_reason = "expert_guard_failed"
            current = previous
            best = previous
            rows.append(
                {
                    "event": "final_cycle_rollback",
                    "cycle": cycle,
                    "rolled_back_to": previous.name,
                    "rolled_back_checkpoint": previous.checkpoint,
                    "failed_checkpoint": guard["target_checkpoint"],
                }
            )
            append_jsonl(SUMMARY_PATH, rows[-1])
            break
        if judgment["worst_win_rate"] <= PASS_THRESHOLD:
            stop_reason = "success_exploiter_threshold"
            break
        if h2h["score_diff_ci95_low"] <= 0.0 <= h2h["score_diff_ci95_high"]:
            stop_reason = "h2h_stagnation"
            break
        worst_protocol = judgment["worst_exploiter"]

    final_candidate = fix_final_candidate(best)
    final_row = {
        "event": "final_cycles_complete",
        "stop_reason": stop_reason,
        "final_candidate": str(final_candidate),
        "final_config": best.config,
        "source_checkpoint": best.checkpoint,
        "elapsed_seconds": time.perf_counter() - started,
    }
    rows.append(final_row)
    append_jsonl(SUMMARY_PATH, final_row)
    write_report(rows, final_row)


def write_league_config(
    cycle: int, target: Target, worst_protocol: str, extra_pool: list[dict[str, Any]]
) -> Path:
    data = yaml.safe_load(Path(LEAGUE_TEMPLATE).read_text(encoding="utf-8")) or {}
    data["base_config"] = target.config
    data["warm_start_checkpoint"] = target.checkpoint
    run = data.setdefault("run", {})
    run["experiment_name"] = f"jax-ppo-final-cycle-c{cycle:02d}"
    run["artifact_root"] = str(ROOT / "league")
    run["seed"] = int(run.get("seed", 20260705)) + 100 + cycle
    league = data.setdefault("league", {})
    league["cycles"] = 1
    league["league_updates_per_cycle"] = 500
    league["snapshot_interval_updates"] = 500
    league["success_exploiter_win_rate"] = PASS_THRESHOLD
    evaluation = data.setdefault("evaluation", {})
    evaluation["games"] = GAMES
    evaluation["batch_games"] = 8192
    exploiter = data.setdefault("exploiter", {})
    spec = exploiter_spec(worst_protocol)
    exploiter["config"] = spec.config
    exploiter["updates"] = load_config(spec.config).run.total_updates
    if spec.resume:
        exploiter["resume"] = spec.resume
    else:
        exploiter.pop("resume", None)
    guards = data.setdefault("guards", {})
    guards["expert_ci_low"] = GUARD_CI_LOW
    guards["max_steps_rate"] = GUARD_MAX_STEPS
    cycle_dir = ROOT / f"cycle_{cycle:02d}"
    tracking = data.setdefault("tracking", {})
    tracking["tracked_summary_path"] = str(cycle_dir / "league_summary.jsonl")
    tracking["report_path"] = str(cycle_dir / "league_report.md")
    anchors = list(data.get("anchors", []))
    names = {item.get("name") for item in anchors}
    for member in extra_pool:
        if member["name"] not in names:
            anchors.append(member)
            names.add(member["name"])
    data["anchors"] = anchors
    path = cycle_dir / "league_config.yaml"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8")
    return path


def evaluate_guard(cycle: int, target: Target) -> dict[str, Any]:
    output = ROOT / f"cycle_{cycle:02d}" / "guard_vs_heuristic_expert.json"
    result = evaluate_checkpoint_vs_static(
        load_config(target.config),
        target.checkpoint,
        "heuristic_expert",
        games=GAMES,
        duplicate=True,
        output=output,
    )
    return {
        "event": "final_cycle_guard_expert",
        "cycle": cycle,
        "target": target.name,
        "target_config": target.config,
        "target_checkpoint": target.checkpoint,
        "json": str(output),
        "passed": result["score_diff_ci95_low"] > GUARD_CI_LOW
        and result["max_steps_rate"] <= GUARD_MAX_STEPS,
        **summary(result),
    }


def evaluate_h2h(cycle: int, target: Target, previous: Target) -> dict[str, Any]:
    output = ROOT / f"cycle_{cycle:02d}" / f"h2h_{target.name}_vs_{previous.name}.json"
    result = evaluate_checkpoint_match(
        load_config(target.config),
        target.checkpoint,
        load_config(previous.config),
        previous.checkpoint,
        games=GAMES,
        duplicate=True,
        output=output,
    )
    return {
        "event": "final_cycle_h2h",
        "cycle": cycle,
        "current": target.name,
        "previous": previous.name,
        "current_checkpoint": target.checkpoint,
        "previous_checkpoint": previous.checkpoint,
        "json": str(output),
        **summary(result),
    }


def run_battery(cycle: int, target: Target) -> list[dict[str, Any]]:
    rows = []
    for spec in EXPLOITERS:
        target_cfg = load_config(target.config)
        exploiter_cfg = load_config(spec.config)
        exploiter_cfg.run.artifact_root = str(ROOT / f"cycle_{cycle:02d}" / "battery")
        exploiter_cfg.run.experiment_name = f"final-c{cycle:02d}-{spec.name}-exploiter"
        train_dir = train_against_checkpoint(
            exploiter_cfg,
            target_cfg,
            target.checkpoint,
            resume=spec.resume,
        )
        output = train_dir / f"eval_vs_{target.name}_duplicate.json"
        result = evaluate_checkpoint_match(
            exploiter_cfg,
            train_dir / "latest",
            target_cfg,
            target.checkpoint,
            games=GAMES,
            duplicate=True,
            output=output,
        )
        rows.append(
            {
                "event": "final_cycle_exploiter",
                "cycle": cycle,
                "exploiter": spec.name,
                "target": target.name,
                "target_config": target.config,
                "target_checkpoint": target.checkpoint,
                "train_dir": str(train_dir),
                "checkpoint": str(train_dir / "latest"),
                "resume": spec.resume,
                "notes": spec.notes,
                "json": str(output),
                **summary(result),
            }
        )
    return rows


def battery_judgment(cycle: int, rows: list[dict[str, Any]]) -> dict[str, Any]:
    worst = max(rows, key=lambda row: row["win_rate"])
    return {
        "event": "final_cycle_battery_judgment",
        "cycle": cycle,
        "worst_exploiter": worst["exploiter"],
        "worst_win_rate": worst["win_rate"],
        "threshold": PASS_THRESHOLD,
        "passed": worst["win_rate"] <= PASS_THRESHOLD,
    }


def exploiter_member_from_league(rows: list[dict[str, Any]], cycle: int) -> dict[str, Any] | None:
    row = latest_row(rows, "exploiter_eval")
    if row is None:
        return None
    checkpoint = Path(row["exploiter_checkpoint"])
    return {
        "name": f"final_cycle_exploiter_c{cycle:02d}",
        "kind": "checkpoint",
        "anchor": False,
        "stalling": False,
        "exploiter": True,
        "config": str(checkpoint.parent / "config.json"),
        "checkpoint": str(checkpoint),
        "recent_win_rate": 1.0 - float(row["win_rate"]),
        "created_cycle": cycle,
        "created_update": 0,
    }


def fix_final_candidate(target: Target) -> Path:
    destination = ROOT / "final_candidate"
    if destination.exists():
        shutil.rmtree(destination)
    shutil.copytree(target.checkpoint, destination)
    (ROOT / "final_candidate_config.txt").write_text(target.config + "\n", encoding="utf-8")
    shutil.copy2(target.config, ROOT / "main_ppo_config.json")
    return destination


def write_report(rows: list[dict[str, Any]], final: dict[str, Any]) -> None:
    guards = [row for row in rows if row.get("event") == "final_cycle_guard_expert"]
    h2h = [row for row in rows if row.get("event") == "final_cycle_h2h"]
    exploiters = [row for row in rows if row.get("event") == "final_cycle_exploiter"]
    judgments = [row for row in rows if row.get("event") == "final_cycle_battery_judgment"]
    lines = [
        f"# Final Cycles and Human Play - {DATE}",
        "",
        "## Part A - Closing Reinforcement Cycles",
        "",
        f"Raw artifacts: `{ROOT}`",
        f"Stop reason: `{final['stop_reason']}`",
        f"Final candidate: `{final['final_candidate']}`",
        f"Final candidate config: `{final['final_config']}`",
        "",
        "### Expert Guard",
        "",
        "| Cycle | Passed | Win rate | Mean diff | CI low | Opened colors | Max-step |",
        "| ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
    ]
    for row in guards:
        lines.append(match_row(row, leading=[str(row["cycle"]), str(row["passed"])]))
    lines.extend(
        [
            "",
            "### Adjacent H2H",
            "",
            "| Cycle | Current | Previous | Win rate | Mean diff | Score CI |",
            "| ---: | --- | --- | ---: | ---: | ---: |",
        ]
    )
    for row in h2h:
        lines.append(
            f"| {row['cycle']} | `{row['current']}` | `{row['previous']}` | "
            f"{row['win_rate']:.4f} | {row['mean_score_diff']:+.4f} | "
            f"[{row['score_diff_ci95_low']:+.4f}, {row['score_diff_ci95_high']:+.4f}] |"
        )
    lines.extend(
        [
            "",
            "### Strengthened Exploiter Battery",
            "",
            "| Cycle | Exploiter | Win rate | Mean diff | CI low | Opened colors | Max-step |",
            "| ---: | --- | ---: | ---: | ---: | ---: | ---: |",
        ]
    )
    for row in exploiters:
        lines.append(match_row(row, leading=[str(row["cycle"]), f"`{row['exploiter']}`"]))
    lines.extend(
        [
            "",
            "### Battery Judgment",
            "",
            "| Cycle | Worst exploiter | Worst win rate | Threshold | Passed |",
            "| ---: | --- | ---: | ---: | ---: |",
        ]
    )
    for row in judgments:
        lines.append(
            f"| {row['cycle']} | `{row['worst_exploiter']}` | {row['worst_win_rate']:.4f} | "
            f"{row['threshold']:.4f} | {row['passed']} |"
        )
    lines.extend(human_play_usage())
    REPORT_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8")


def human_play_usage() -> list[str]:
    return [
        "",
        "## Part B - Human Play Interface",
        "",
        "Start a single game:",
        "",
        "```bash",
        "uv run --with 'jax[cuda12]' lost-cities-jax-ppo play \\",
        "  --checkpoint /mnt/2tbhdd/coolrl-lost-cities-artifacts/final-cycles/2026-07-05/final_candidate \\",
        "  --seat 0",
        "```",
        "",
        "Start a duplicate set with one shared shuffle and swapped seats:",
        "",
        "```bash",
        "uv run --with 'jax[cuda12]' lost-cities-jax-ppo play \\",
        "  --checkpoint /mnt/2tbhdd/coolrl-lost-cities-artifacts/final-cycles/2026-07-05/final_candidate \\",
        "  --seat 0 --duplicate",
        "```",
        "",
        "Summarize logged human games:",
        "",
        "```bash",
        "uv run lost-cities-jax-ppo human-play summarize \\",
        "  --log-dir /mnt/2tbhdd/coolrl-lost-cities-artifacts/human-play/",
        "```",
        "",
        "Move syntax: `play R7 draw deck`, `discard G3 draw Y`, or `play RHS draw deck`.",
        "The renderer shows only the human hand, both boards, all public discard piles, deck count, and current board score differential. Opponent hand and deck order are not rendered.",
        "Every game is appended to `/mnt/2tbhdd/coolrl-lost-cities-artifacts/human-play/games.jsonl` with deck seed/index, full action list, AI top-3 policy actions/probabilities, value outputs, scoring breakdown, and optional human comment.",
    ]


def match_row(row: dict[str, Any], leading: list[str]) -> str:
    return (
        "| " + " | ".join(leading) + f" | {row['win_rate']:.4f} | {row['mean_score_diff']:+.4f} | "
        f"{row['score_diff_ci95_low']:+.4f} | {row['opened_colors_per_game']:.4f} | "
        f"{row['max_steps_rate']:.4f} |"
    )


def summary(result: dict[str, Any]) -> dict[str, Any]:
    keys = [
        "games",
        "wins",
        "losses",
        "ties",
        "win_rate",
        "wilson_low",
        "wilson_high",
        "mean_score_diff",
        "score_diff_ci95_low",
        "score_diff_ci95_high",
        "mean_game_length",
        "max_steps_rate",
        "opened_colors_per_game",
        "play_action_rate",
        "positive_expeditions_per_game",
    ]
    return {key: result[key] for key in keys if key in result}


def exploiter_spec(name: str) -> ExploiterSpec:
    for spec in EXPLOITERS:
        if spec.name == name:
            return spec
    raise ValueError(f"unknown exploiter protocol: {name}")


def latest_row(rows: list[dict[str, Any]], event: str) -> dict[str, Any] | None:
    return next((row for row in reversed(rows) if row.get("event") == event), None)


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def append_jsonl(path: Path, row: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(row, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()
