"""Scripted Snake agents: the oracle used for supervision and the baselines.

All agents implement `act(state: SnakeState) -> str`.

Path search is *time-aware*: a body cell at index i (head=0) of a snake of
length L becomes free after L - i moves, so a path may pass where the tail
currently is if it arrives late enough. This makes BFS paths exact as long as
no food is eaten along the way (food is only ever the final cell).
"""
from __future__ import annotations

from collections import deque
from functools import lru_cache

import numpy as np

from .env import ACTIONS, DELTAS, SnakeState


@lru_cache(maxsize=16)
def _neighbors(h: int, w: int) -> tuple[tuple[tuple[int, str], ...], ...]:
    out = []
    for cell in range(h * w):
        r, c = divmod(cell, w)
        nb = []
        for a in ACTIONS:
            dr, dc = DELTAS[a]
            nr, nc = r + dr, c + dc
            if 0 <= nr < h and 0 <= nc < w:
                nb.append((nr * w + nc, a))
        out.append(tuple(nb))
    return tuple(out)


def _flat(state: SnakeState) -> tuple[list[int], int]:
    w = state.width
    body = [r * w + c for r, c in state.body]
    food = state.food[0] * w + state.food[1] if state.food is not None else -1
    return body, food


def _free_times(body: list[int], area: int) -> list[int]:
    ft = [0] * area
    L = len(body)
    for i, cell in enumerate(body):
        ft[cell] = L - i
    return ft


def _bfs(nbrs, area: int, start: int, ft: list[int], goal: int = -1):
    dist = [-1] * area
    parent = [-1] * area
    dist[start] = 0
    q = deque([start])
    while q:
        u = q.popleft()
        du = dist[u] + 1
        for v, _ in nbrs[u]:
            if dist[v] < 0 and du >= ft[v]:
                dist[v] = du
                parent[v] = u
                if v == goal:
                    return dist, parent
                q.append(v)
    return dist, parent


def _reachable_count(nbrs, area: int, start: int, ft: list[int]) -> int:
    dist, _ = _bfs(nbrs, area, start, ft)
    return sum(1 for d in dist if d > 0)


def _tail_reachable(nbrs, area: int, body: list[int]) -> tuple[bool, int]:
    if len(body) < 2:
        return True, 0
    ft = _free_times(body, area)
    tail = body[-1]
    dist, _ = _bfs(nbrs, area, body[0], ft, goal=tail)
    return dist[tail] > 0, dist[tail]


def safe_moves(state: SnakeState) -> list[str]:
    """Moves that do not die on the very next step."""
    body, _ = _flat(state)
    nbrs = _neighbors(state.height, state.width)
    ft = _free_times(body, state.height * state.width)
    return [a for v, a in nbrs[body[0]] if ft[v] <= 1]


class Agent:
    name = "agent"

    def act(self, state: SnakeState) -> str:  # pragma: no cover - interface
        raise NotImplementedError

    def reset(self, seed: int | None = None) -> None:
        pass


class RandomAgent(Agent):
    name = "random"

    def __init__(self, seed: int = 0):
        self.rng = np.random.default_rng(seed)

    def reset(self, seed=None):
        if seed is not None:
            self.rng = np.random.default_rng(seed)

    def act(self, state):
        return ACTIONS[int(self.rng.integers(4))]


class GreedyFoodAgent(Agent):
    """Minimise Manhattan distance to food; avoid *immediately* fatal moves when possible."""
    name = "greedy"

    def act(self, state):
        safe = safe_moves(state) or list(ACTIONS)
        hr, hc = state.body[0]
        fr, fc = state.food if state.food else (hr, hc)
        return min(safe, key=lambda a: abs(hr + DELTAS[a][0] - fr) + abs(hc + DELTAS[a][1] - fc))


class BFSAgent(Agent):
    """Shortest (time-aware) path to food, no look-ahead; else the safe move with most space."""
    name = "bfs"

    def act(self, state):
        h, w = state.height, state.width
        area = h * w
        nbrs = _neighbors(h, w)
        body, food = _flat(state)
        ft = _free_times(body, area)
        if food >= 0:
            path = _path_to(nbrs, area, body[0], ft, food)
            if path:
                return _direction(w, body[0], path[0])
        return _max_space_move(nbrs, area, body) or ACTIONS[0]


class SurvivalAgent(Agent):
    """Flood-fill heuristic: maximise reachable space, break ties by distance to food."""
    name = "survival"

    def act(self, state):
        h, w = state.height, state.width
        area = h * w
        nbrs = _neighbors(h, w)
        body, food = _flat(state)
        ft = _free_times(body, area)
        best, best_key = None, None
        fr, fc = divmod(food, w) if food >= 0 else divmod(body[0], w)
        for v, a in nbrs[body[0]]:
            if ft[v] > 1:
                continue
            nb = [v] + (body if v == food else body[:-1])
            space = _reachable_count(nbrs, area, v, _free_times(nb, area))
            r, c = divmod(v, w)
            key = (min(space, len(nb) + 1), -(abs(r - fr) + abs(c - fc)))
            if best_key is None or key > best_key:
                best, best_key = a, key
        return best or ACTIONS[0]


class OracleAgent(Agent):
    """Strong teacher: eat only if the snake can still reach its tail afterwards,
    otherwise chase the tail along the longest safe route, otherwise maximise space.

    `analyze` also returns the *set* of equally good moves. Ties are broken by a fixed
    priority (UP > DOWN > LEFT > RIGHT) instead of BFS queue order, so the hard label
    is a learnable function of the state (and this rule also scored best: 95.0 vs
    91.5-92.1 for "close the larger/smaller axis gap first", 60 games 10x10).
    """
    name = "oracle"

    def act(self, state):
        return self.analyze(state)[0]

    def analyze(self, state) -> tuple[str, list[str], str]:
        """-> (move, optimal_moves, mode) with mode in {eat, tail, space, dead}."""
        h, w = state.height, state.width
        area = h * w
        nbrs = _neighbors(h, w)
        body, food = _flat(state)
        L = len(body)
        ft = _free_times(body, area)
        head = body[0]
        safe = [(v, a) for v, a in nbrs[head] if ft[v] <= 1]
        if not safe:
            return ACTIONS[0], list(ACTIONS), "dead"
        if food >= 0:
            dist, _ = _bfs(nbrs, area, head, ft, goal=food)
            d = dist[food]
            if d > 0:
                good = []
                for v, a in safe:
                    if v == food:
                        virtual = [v] + body
                    else:
                        nb = [v] + body[:-1]
                        path = _path_to(nbrs, area, v, _free_times(nb, area), food)
                        if len(path) != d - 1:
                            continue
                        virtual = (path[::-1] + nb)[: L + 1]
                    if len(virtual) >= area or _tail_reachable(nbrs, area, virtual)[0]:
                        good.append(a)
                if good:
                    return good[0], good, "eat"
        # tail chasing: keep an escape route, prefer the move farthest from the tail
        cands = []
        for v, a in safe:
            nb = [v] + (body if v == food else body[:-1])
            ok, dt = _tail_reachable(nbrs, area, nb)
            if ok:
                cands.append((dt, a))
        if cands:
            best = max(dt for dt, _ in cands)
            good = [a for dt, a in cands if dt == best]
            return good[0], good, "tail"
        spaces = []
        for v, a in safe:
            nb = [v] + (body if v == food else body[:-1])
            spaces.append((_reachable_count(nbrs, area, v, _free_times(nb, area)), a))
        best = max(sp for sp, _ in spaces)
        good = [a for sp, a in spaces if sp == best]
        return good[0], good, "space"


def _path_to(nbrs, area, start, ft, goal) -> list[int]:
    dist, parent = _bfs(nbrs, area, start, ft, goal)
    if dist[goal] <= 0:
        return []
    path = [goal]
    while parent[path[-1]] != start:
        path.append(parent[path[-1]])
    return path[::-1]


def _direction(w: int, a: int, b: int) -> str:
    ar, ac = divmod(a, w)
    br, bc = divmod(b, w)
    return {(-1, 0): "UP", (1, 0): "DOWN", (0, -1): "LEFT", (0, 1): "RIGHT"}[(br - ar, bc - ac)]


def _max_space_move(nbrs, area, body) -> str | None:
    ft = _free_times(body, area)
    best, best_space = None, -1
    for v, a in nbrs[body[0]]:
        if ft[v] > 1:
            continue
        nb = [v] + body[:-1]
        space = _reachable_count(nbrs, area, v, _free_times(nb, area))
        if space > best_space:
            best, best_space = a, space
    return best


AGENTS = {c.name: c for c in (RandomAgent, GreedyFoodAgent, BFSAgent, SurvivalAgent, OracleAgent)}


def make_agent(name: str, seed: int = 0) -> Agent:
    cls = AGENTS[name]
    return cls(seed) if cls is RandomAgent else cls()