"""Compact adversarial rollout for multiplayer Battlesnake positions."""

from __future__ import annotations

from itertools import product
from time import perf_counter

from snakes.bitboard import BitBoard

Body = tuple[int, ...]
EnemyBodies = tuple[Body, ...]
StateKey = tuple[Body, EnemyBodies, int, int]

class CompactSurvivalSearch:
  """Small paranoid beam search with simultaneous enemy responses.

  It is deliberately narrower than full multiplayer minimax: each enemy keeps
  only its most dangerous replies and the combined response beam is capped.
  This models moving opponents without exhausting the request deadline.
  """

  DEATH = -5000.0

  def __init__(
    self,
    board: BitBoard,
    food: set[tuple[int, int]],
    is_constrictor: bool,
    deadline: float | None,
    branch: int,
    enemy_branch: int = 2,
    response_cap: int = 8,
  ) -> None:
    self.board = board
    self.food_bits = board.set_to_bits(food)
    self.is_constrictor = is_constrictor
    self.deadline = deadline
    self.branch = max(1, branch)
    self.enemy_branch = max(1, enemy_branch)
    self.response_cap = max(1, response_cap)
    self.cache: dict[StateKey, float] = {}
    self.body_bits_cache: dict[Body, int] = {}
    self.nodes = 0
    self.cache_hits = 0

  def body_from_dicts(self, body: list[dict]) -> Body:
    return tuple(self.board.idx(segment["x"], segment["y"]) for segment in body)

  def search_selected(
    self,
    my_body: list[dict],
    enemies: list[dict],
    target: tuple[int, int],
    depth: int,
  ) -> float:
    mine = self.body_from_dicts(my_body)
    enemy_bodies = tuple(self.body_from_dicts(enemy["body"]) for enemy in enemies)
    target_idx = self.board.idx(*target)
    if not self.board.neighbors_of(mine[0]) & (1 << target_idx):
      return self.DEATH
    return self._selected_root(mine, enemy_bodies, self.food_bits, target_idx, depth)

  def _selected_root(
    self, mine: Body, enemies: EnemyBodies, food_bits: int, target: int, depth: int,
  ) -> float:
    replies = self._enemy_responses(enemies, mine, target, food_bits)
    if not replies:
      replies = [()]
    worst = float("inf")
    for response in replies:
      if self._out_of_time():
        break
      child = self._advance(mine, enemies, target, response, food_bits)
      if child is None:
        value = self.DEATH
      else:
        next_mine, next_enemies, next_food = child
        value = self._evaluate(next_mine, next_enemies)
        if depth > 1 and value > self.DEATH:
          value += self._search(next_mine, next_enemies, next_food, depth - 1) * 0.72
      worst = min(worst, value)
    return self._evaluate(mine, enemies) if worst == float("inf") else worst

  def _search(self, mine: Body, enemies: EnemyBodies, food_bits: int, depth: int) -> float:
    self.nodes += 1
    if self._out_of_time() or depth <= 0:
      return 0.0
    key = (mine, enemies, food_bits, depth)
    cached = self.cache.get(key)
    if cached is not None:
      self.cache_hits += 1
      return cached

    occupied = self._occupied(mine, enemies)
    targets = list(self._iter_bits(self.board.neighbors_of(mine[0])))
    ranked: list[tuple[float, int]] = []
    for target in targets:
      # Collision legality is finalized simultaneously because eating controls
      # whether tails vacate.
      ate = bool((1 << target) & food_bits)
      own_tail_blocked = self.is_constrictor or ate
      body_blocked = self._body_bits(mine if own_tail_blocked else mine[:-1])
      enemy_blocked = 0
      for enemy in enemies:
        enemy_blocked |= self._body_bits(enemy[:-1] if not self.is_constrictor else enemy)
      if (1 << target) & (body_blocked | enemy_blocked):
        continue
      free_space = self.board.flood_count(target, occupied & ~(1 << target))
      ranked.append((free_space + (20 if ate else 0), target))
    ranked.sort(reverse=True)

    if not ranked:
      return self.DEATH

    best = self.DEATH
    for _, target in ranked[:self.branch]:
      replies = self._enemy_responses(enemies, mine, target, food_bits) or [()]
      worst = float("inf")
      for response in replies:
        if self._out_of_time():
          break
        child = self._advance(mine, enemies, target, response, food_bits)
        if child is None:
          value = self.DEATH
        else:
          next_mine, next_enemies, next_food = child
          value = self._evaluate(next_mine, next_enemies)
          if depth > 1 and value > self.DEATH:
            value += self._search(next_mine, next_enemies, next_food, depth - 1) * 0.72
        worst = min(worst, value)
      if worst != float("inf"):
        best = max(best, worst)

    if not self._out_of_time() and len(self.cache) < 16_384:
      self.cache[key] = best
    return best

  def _enemy_responses(
    self, enemies: EnemyBodies, mine: Body, my_target: int, food_bits: int,
  ) -> list[tuple[int, ...]]:
    if not enemies:
      return []
    choices: list[list[int]] = []
    my_length_after = len(mine) + int(bool((1 << my_target) & food_bits))
    for enemy in enemies:
      ranked: list[tuple[float, int]] = []
      for target in self._iter_bits(self.board.neighbors_of(enemy[0])):
        ate = bool((1 << target) & food_bits)
        enemy_length_after = len(enemy) + int(ate)
        score = 0.0
        if target == my_target:
          score += 1000.0 if enemy_length_after >= my_length_after else -1000.0
        tx, ty = self.board.coord(target)
        mx, my = self.board.coord(my_target)
        score -= abs(tx - mx) + abs(ty - my)
        score += self.board.open_neighbor_count(target, self._occupied(mine, enemies)) * 3.0
        score += 20.0 if ate else 0.0
        ranked.append((score, target))
      ranked.sort(reverse=True)
      choices.append([target for _, target in ranked[:self.enemy_branch]])

    responses: list[tuple[int, ...]] = []
    for response in product(*choices):
      responses.append(response)
      if len(responses) >= self.response_cap:
        break
    return responses

  def _advance(
    self,
    mine: Body,
    enemies: EnemyBodies,
    my_target: int,
    enemy_targets: tuple[int, ...],
    food_bits: int,
  ) -> tuple[Body, EnemyBodies, int] | None:
    my_ate = bool((1 << my_target) & food_bits)
    next_mine = self._advance_body(mine, my_target, my_ate)
    next_enemies = tuple(
      self._advance_body(body, target, bool((1 << target) & food_bits))
      for body, target in zip(enemies, enemy_targets)
    )

    # Body and self collisions after all tails have moved.
    if my_target in next_mine[1:]:
      return None
    if any(my_target in enemy[1:] for enemy in next_enemies):
      return None

    surviving: list[Body] = []
    for index, enemy in enumerate(next_enemies):
      target = enemy[0]
      dead = target in enemy[1:] or target in next_mine[1:]
      dead = dead or any(
        target in other[1:] for other_index, other in enumerate(next_enemies)
        if other_index != index
      )
      if target == my_target:
        if len(enemy) >= len(next_mine):
          return None
        dead = True
      if not dead:
        # Enemy/enemy head collisions remove equal-length snakes and the shorter.
        for other_index, other in enumerate(next_enemies):
          if other_index != index and target == other[0] and len(enemy) <= len(other):
            dead = True
            break
      if not dead:
        surviving.append(enemy)

    eaten = (1 << my_target) if my_ate else 0
    for body, target in zip(enemies, enemy_targets):
      if (1 << target) & food_bits:
        eaten |= 1 << target
    return next_mine, tuple(surviving), food_bits & ~eaten

  def _evaluate(self, mine: Body, enemies: EnemyBodies) -> float:
    blocked = self._occupied(mine, enemies) & ~(1 << mine[0])
    space = self.board.flood_count(mine[0], blocked)
    liberties = self.board.open_neighbor_count(mine[0], blocked)
    if liberties == 0 or space < len(mine):
      return self.DEATH
    enemy_pressure = 0.0
    for enemy in enemies:
      enemy_blocked = blocked & ~(1 << enemy[0])
      enemy_space = self.board.flood_count(enemy[0], enemy_blocked)
      enemy_liberties = self.board.open_neighbor_count(enemy[0], enemy_blocked)
      enemy_pressure += max(0, 3 - enemy_liberties) * 18.0
      if len(mine) > len(enemy):
        enemy_pressure += max(0, 8 - enemy_space) * 8.0
    return space * 1.9 + liberties * 32.0 + enemy_pressure - len(enemies) * 4.0

  def _occupied(self, mine: Body, enemies: EnemyBodies) -> int:
    occupied = self._body_bits(mine)
    for enemy in enemies:
      occupied |= self._body_bits(enemy)
    return occupied

  def _body_bits(self, body: Body) -> int:
    cached = self.body_bits_cache.get(body)
    if cached is not None:
      return cached
    bits = 0
    for cell in body:
      bits |= 1 << cell
    if len(self.body_bits_cache) < 16_384:
      self.body_bits_cache[body] = bits
    return bits

  def _advance_body(self, body: Body, target: int, ate: bool) -> Body:
    if self.is_constrictor or ate:
      return (target,) + body
    return (target,) + body[:-1]

  @staticmethod
  def _iter_bits(bits: int):
    while bits:
      bit = bits & -bits
      yield bit.bit_length() - 1
      bits ^= bit

  def _out_of_time(self) -> bool:
    return self.deadline is not None and perf_counter() >= self.deadline