"""Play full Snake games with any policy and report game-level metrics."""
from __future__ import annotations

import time
from typing import Callable, Sequence

import numpy as np

from .agents import Agent, make_agent, safe_moves
from .env import ACTIONS, SnakeConfig, SnakeEnv, SnakeState

BatchPolicy = Callable[[Sequence[SnakeState]], Sequence[str]]


def agent_policy(agent: Agent) -> BatchPolicy:
    return lambda states: [agent.act(s) for s in states]


def model_policy(decider, question: str = "Choose the next move", choices=None) -> BatchPolicy:
    """Greedy MiniJEV policy; choices default to the model's own labels (legacy "UP" or "move up")."""
    task = decider.task
    choices = list(choices) if choices is not None else list(task.default_choices)
    actions = [task.action_of(c) for c in choices]

    def policy(states):
        probs = decider.probs(states, question, choices)
        return [actions[i] for i in probs.argmax(1)]
    return policy


def play_games(policy: BatchPolicy, n_games: int = 100, env_cfg: dict | None = None, seed: int = 12345,
               parallel: int = 64, max_decisions: int | None = None) -> dict:
    """Run `n_games` complete games (`parallel` at a time, batched decisions).

    Game seeds are fixed (seed + i), so all agents are compared on identical games.
    """
    cfg = SnakeConfig.from_dict(env_cfg)
    results = []
    next_game = 0
    active: list[SnakeEnv] = []
    fatal = decisions = 0
    t0 = time.perf_counter()
    policy_time = 0.0

    def spawn():
        nonlocal next_game
        env = SnakeEnv(cfg, seed=seed + next_game)
        next_game += 1
        return env

    while next_game < min(parallel, n_games):
        active.append(spawn())
    while active:
        states = [e.state() for e in active]
        tp = time.perf_counter()
        actions = policy(states)
        policy_time += time.perf_counter() - tp
        still = []
        for env, s, a in zip(active, states, actions):
            decisions += 1
            if a not in safe_moves(s):
                fatal += 1
            env.step(a)
            if env.done:
                results.append({"score": env.score, "steps": env.steps, "end": env.end_reason})
                if next_game < n_games:
                    still.append(spawn())
            else:
                still.append(env)
        active = still
        if max_decisions and decisions >= max_decisions:
            break
    wall = time.perf_counter() - t0
    scores = np.array([r["score"] for r in results]) if results else np.zeros(1)
    steps = np.array([r["steps"] for r in results]) if results else np.zeros(1)
    ends: dict[str, int] = {}
    for r in results:
        ends[r["end"]] = ends.get(r["end"], 0) + 1
    n = max(1, len(results))
    return {
        "games": len(results),
        "board": f"{cfg.height}x{cfg.width}",
        "avg_score": float(scores.mean()), "median_score": float(np.median(scores)),
        "max_score": int(scores.max()), "min_score": int(scores.min()),
        "avg_survival_steps": float(steps.mean()), "food_collected": int(scores.sum()),
        "max_possible_score": cfg.height * cfg.width - cfg.initial_length,
        "win_rate": ends.get("win", 0) / n,
        "death_rate": (ends.get("wall", 0) + ends.get("body", 0)) / n,  # collisions
        "starvation_rate": ends.get("starvation", 0) / n,
        "end_reasons": ends,
        "decisions": decisions,
        "fatal_move_rate": fatal / max(1, decisions),  # chose a move that dies immediately
        "policy_seconds": round(policy_time, 3),
        "wall_seconds": round(wall, 3),
    }


def baseline_suite(n_games: int = 100, env_cfg: dict | None = None, seed: int = 12345,
                   agents: Sequence[str] = ("random", "greedy", "bfs", "survival", "oracle")) -> dict:
    return {name: play_games(agent_policy(make_agent(name, seed)), n_games, env_cfg, seed, parallel=1)
            for name in agents}