"""Headless, deterministic Snake.

Cells are flattened indices `r * width + c`. The body is stored head-first in a
deque plus a bytearray occupancy grid, so a step is O(1) and thousands of games
run without any rendering.
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field

import numpy as np

ACTIONS = ("UP", "DOWN", "LEFT", "RIGHT")
DELTAS = {"UP": (-1, 0), "DOWN": (1, 0), "LEFT": (0, -1), "RIGHT": (0, 1)}
OPPOSITE = {"UP": "DOWN", "DOWN": "UP", "LEFT": "RIGHT", "RIGHT": "LEFT"}


@dataclass
class RewardConfig:
    food: float = 1.0
    death: float = -1.0
    step: float = 0.0  # small survival bonus; keep tiny, starvation limit prevents stalling exploits
    win: float = 1.0

    @classmethod
    def from_dict(cls, d: dict | None) -> "RewardConfig":
        return cls(**(d or {}))


@dataclass
class SnakeConfig:
    height: int = 10
    width: int = 10
    initial_length: int = 3
    # "die": reversing into the neck is a collision (every choice has a distinct outcome)
    # "ignore": reversal keeps the current heading (classic arcade behaviour)
    reverse_action: str = "die"
    starvation_steps: int | None = None  # None -> 2 * area
    max_steps: int | None = None
    reward: RewardConfig = field(default_factory=RewardConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "SnakeConfig":
        d = dict(d or {})
        board = d.pop("board_size", None)
        if board is not None:
            d["height"] = d["width"] = int(board)
        d["reward"] = RewardConfig.from_dict(d.get("reward"))
        return cls(**d)

    @property
    def starvation_limit(self) -> int:
        return self.starvation_steps or 2 * self.height * self.width


@dataclass
class SnakeState:
    """Immutable-ish snapshot, the `state` MiniJEV receives for Snake."""
    height: int
    width: int
    body: list[tuple[int, int]]  # head first
    food: tuple[int, int] | None
    direction: str
    steps_since_food: int = 0

    def to_dict(self) -> dict:
        return {"height": self.height, "width": self.width, "body": [list(p) for p in self.body],
                "food": list(self.food) if self.food else None, "direction": self.direction,
                "steps_since_food": self.steps_since_food}

    @classmethod
    def from_dict(cls, d: dict) -> "SnakeState":
        return cls(d["height"], d["width"], [tuple(p) for p in d["body"]],
                   tuple(d["food"]) if d.get("food") else None, d["direction"], d.get("steps_since_food", 0))


class SnakeEnv:
    def __init__(self, config: SnakeConfig | None = None, seed: int | None = None):
        self.cfg = config or SnakeConfig()
        self.h, self.w = self.cfg.height, self.cfg.width
        self.area = self.h * self.w
        self.rng = np.random.default_rng(seed)
        self.body: deque[int] = deque()
        self.occ = bytearray(self.area)
        self.reset()

    # ------------------------------------------------------------------ core
    def seed(self, seed: int) -> None:
        self.rng = np.random.default_rng(seed)

    def reset(self, seed: int | None = None) -> SnakeState:
        if seed is not None:
            self.seed(seed)
        self.body.clear()
        self.occ = bytearray(self.area)
        self.direction = ACTIONS[int(self.rng.integers(4))]
        dr, dc = DELTAS[self.direction]
        L = min(self.cfg.initial_length, min(self.h, self.w) // 2)
        # place head so the body (trailing behind it) fits on the board
        r = int(self.rng.integers(L, self.h - L)) if self.h > 2 * L else self.h // 2
        c = int(self.rng.integers(L, self.w - L)) if self.w > 2 * L else self.w // 2
        for i in range(L):
            cell = (r - dr * i) * self.w + (c - dc * i)
            self.body.append(cell)
            self.occ[cell] = 1
        self.score = 0
        self.steps = 0
        self.steps_since_food = 0
        self.done = False
        self.end_reason: str | None = None
        self._place_food()
        return self.state()

    def _place_food(self) -> None:
        free = self.area - len(self.body)
        if free <= 0:
            self.food = -1
            return
        # k-th free cell, found by scanning; cheap versus rejection sampling on full boards
        k = int(self.rng.integers(free))
        if free > self.area // 4:
            while True:
                cell = int(self.rng.integers(self.area))
                if not self.occ[cell]:
                    self.food = cell
                    return
        for cell in range(self.area):
            if not self.occ[cell]:
                if k == 0:
                    self.food = cell
                    return
                k -= 1

    def step(self, action: str | int) -> tuple[float, bool]:
        """Returns (reward, done). Call `state()` for the observation."""
        if self.done:
            raise RuntimeError("step() on finished game; call reset()")
        if not isinstance(action, str):
            action = ACTIONS[int(action)]
        if self.cfg.reverse_action == "ignore" and len(self.body) > 1 and action == OPPOSITE[self.direction]:
            action = self.direction
        rw = self.cfg.reward
        dr, dc = DELTAS[action]
        head = self.body[0]
        r, c = divmod(head, self.w)
        nr, nc = r + dr, c + dc
        self.direction = action
        self.steps += 1
        self.steps_since_food += 1
        if not (0 <= nr < self.h and 0 <= nc < self.w):
            return self._finish("wall", rw.death)
        cell = nr * self.w + nc
        eating = cell == self.food
        tail = self.body[-1]
        if self.occ[cell] and not (cell == tail and not eating):
            return self._finish("body", rw.death)
        if not eating:
            self.body.pop()
            self.occ[tail] = 0
        self.body.appendleft(cell)
        self.occ[cell] = 1
        reward = rw.step
        if eating:
            self.score += 1
            self.steps_since_food = 0
            reward += rw.food
            if len(self.body) == self.area:
                return self._finish("win", reward + rw.win)
            self._place_food()
        if self.steps_since_food >= self.cfg.starvation_limit:
            return self._finish("starvation", rw.death)
        if self.cfg.max_steps and self.steps >= self.cfg.max_steps:
            self.done, self.end_reason = True, "max_steps"
            return reward, True
        return reward, False

    def _finish(self, reason: str, reward: float) -> tuple[float, bool]:
        self.done, self.end_reason = True, reason
        return reward, True

    # --------------------------------------------------------------- views
    def state(self) -> SnakeState:
        w = self.w
        return SnakeState(self.h, self.w, [divmod(b, w) for b in self.body],
                          divmod(self.food, w) if self.food >= 0 else None,
                          self.direction, self.steps_since_food)

    def render_ascii(self) -> str:
        grid = [["." for _ in range(self.w)] for _ in range(self.h)]
        for i, b in enumerate(self.body):
            r, c = divmod(b, self.w)
            grid[r][c] = "H" if i == 0 else "o"
        if self.food >= 0:
            r, c = divmod(self.food, self.w)
            grid[r][c] = "*"
        return "\n".join("".join(row) for row in grid)


class VectorSnakeEnv:
    """N independent games stepped together (auto-reset); used for batched
    inference in RL/evaluation. Deterministic given `seed`."""

    def __init__(self, n: int, config: SnakeConfig | None = None, seed: int = 0):
        self.envs = [SnakeEnv(config, seed=seed + i) for i in range(n)]
        self.next_seed = seed + n
        self.n = n

    def states(self) -> list[SnakeState]:
        return [e.state() for e in self.envs]

    def step(self, actions) -> tuple[np.ndarray, np.ndarray, list[dict]]:
        rewards = np.zeros(self.n, dtype=np.float32)
        dones = np.zeros(self.n, dtype=bool)
        infos: list[dict] = [{} for _ in range(self.n)]
        for i, (env, a) in enumerate(zip(self.envs, actions)):
            rewards[i], dones[i] = env.step(a)
            if dones[i]:
                infos[i] = {"score": env.score, "steps": env.steps, "end_reason": env.end_reason}
                env.reset(seed=self.next_seed)
                self.next_seed += 1
        return rewards, dones, infos