from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from pathlib import Path
from time import perf_counter

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from tower_game.ai import LocalOpenAIGateway, MockGateway, load_model_config
from tower_game.engine import TowerGame


def prepared_state(game: TowerGame, seed: int = 7001):
    state, _ = game.new_character(seed=seed)
    game.confirm_allocation(state, [6, 6, 6, 6])
    game.prepare_run(state)
    game.choose_starter_skill(state, 0)
    return state


def record(rows, operation, gateway, started, fallback=False):
    metrics = dict(gateway.agent.last_metrics)
    rows.append(
        {
            "operation": operation,
            "latency_ms": round((perf_counter() - started) * 1000),
            "fallback": fallback,
            "last_error": getattr(gateway, "_last_error", ""),
            **metrics,
        }
    )


def repeated_verification(gateway: LocalOpenAIGateway, repetitions: int):
    game = TowerGame(gateway)
    state = prepared_state(game)
    state["current_floor"] = 10
    state["boss_deck"] = [
        move.model_dump()
        for move in MockGateway().boss_deck(state, final=True).moves
    ]
    rows = []
    for index in range(repetitions):
        operations = [
            ("run_setup", lambda: gateway.run_setup(state)),
            ("class_evolution", lambda: gateway.evolution(state, 1 + index % 2)),
            (
                "boss_package",
                lambda: gateway.boss_package(
                    state, game.assets.ids("boss"), final=bool(index % 2)
                ),
            ),
            ("boss_adjustment", lambda: gateway.boss_adjustment(state)),
            ("boss_turn_decision", lambda: gateway.boss_turn_decision(state)),
            ("ascension", lambda: gateway.ascension(state)),
        ]
        for name, operation in operations:
            started = perf_counter()
            operation()
            record(rows, name, gateway, started, bool(getattr(gateway, "_last_error", "")))
    return rows


def full_run(gateway: LocalOpenAIGateway):
    print("AI Ascension: generating run setup...", flush=True)
    game = TowerGame(gateway)
    state = prepared_state(game, seed=8101)
    guard = 0
    last_phase = ""
    while state["game_phase"] != "ascension" and guard < 200:
        guard += 1
        phase = state["game_phase"]
        if phase != last_phase:
            print(
                f"AI Ascension: {phase} / floor {state.get('current_floor')}",
                flush=True,
            )
            last_phase = phase
        if phase == "combat":
            state["enemy_evasion"] = 0
            state["enemy_hp"] = 1
            game.act(state, "strike")
            if state.get("boss_thinking") and state["game_phase"] == "combat":
                decision, status = game.generate_boss_decision(state)
                game.apply_boss_decision(state, decision, status)
        elif phase == "victory":
            if state["victory_step"] == "level_skill":
                game.choose_level_skill(state, 0)
                if state["game_phase"] == "skill_replacement":
                    game.replace_skill(state, 0)
            elif state["victory_step"] == "loot":
                while not state["proceed_ready"]:
                    index = next(
                        i
                        for i, reward in enumerate(state["pending_loot"])
                        if not reward.get("claimed")
                    )
                    game.choose_loot(state, index)
                game.proceed(state)
            else:
                game.proceed(state)
        elif phase == "evolution_loading":
            evolution, status = game.generate_evolution(state)
            game.complete_evolution_generation(state, evolution, status)
        elif phase == "evolution_reveal":
            game.embrace_evolution(state)
        elif phase == "evolution_healing":
            game.heal_choice(state, True)
        elif phase in {"skill_replacement", "evolution_skill_replace"}:
            game.replace_skill(state, 0)
        elif phase == "boss_loading":
            package, decision, status = game.generate_boss_intro(state)
            game.complete_boss_intro(state, package, decision, status)
        elif phase == "ascension_loading":
            passive, status = game.generate_ascension(state)
            game.complete_ascension(state, passive, status)
        elif phase == "defeat":
            raise RuntimeError("AI-enabled verification run was defeated")
        else:
            raise RuntimeError(f"unexpected phase during verification: {phase}")
    if state["game_phase"] != "ascension":
        raise RuntimeError("AI-enabled run did not reach Ascension")
    return {
        "floor": state["current_floor"],
        "ascension": state["ascension_level"],
        "fallbacks": state["agent_status"]["fallback_count"],
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--repetitions", type=int, default=20)
    parser.add_argument("--skip-full-run", action="store_true")
    parser.add_argument(
        "--output", type=Path, default=Path(".local/llama.cpp/verification.json")
    )
    args = parser.parse_args()

    config = load_model_config()
    config["backend"] = "local_openai"
    gateway = LocalOpenAIGateway(config)
    rows = repeated_verification(gateway, args.repetitions)
    summary = {
        "calls": len(rows),
        "fallbacks": sum(1 for row in rows if row["fallback"]),
        "operations": dict(Counter(row["operation"] for row in rows)),
        "average_latency_ms": round(
            sum(row["latency_ms"] for row in rows) / max(1, len(rows))
        ),
        "results": rows,
    }
    if not args.skip_full_run:
        summary["full_run"] = full_run(gateway)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps({key: value for key, value in summary.items() if key != "results"}, indent=2))


if __name__ == "__main__":
    main()