"""Deadline-aware simultaneous duel search using compact tuple bodies and bitboards."""

from __future__ import annotations

from dataclasses import dataclass
from time import perf_counter
from typing import Iterable

from snakes.bitboard import BitBoard

Body = tuple[int, ...]

@dataclass(frozen=True, slots=True)
class DuelState:
  my_body: Body
  enemy_body: Body
  food_bits: int
  my_health: int
  enemy_health: int
  previous_hazard_bits: int

class BitboardDuelSearch:
  """Iterative-deepening paranoid minimax for a two-snake game.

  The public API still accepts Battlesnake body dictionaries. Search nodes use
  flat cell indices, immutable tuples, and integer masks to avoid allocation of
  coordinate dictionaries and sets in the hot path.
  """

  WIN = 100_000.0
  LOSS = -100_000.0

  def __init__(
    self,
    board: BitBoard,
    food: Iterable[tuple[int, int]],
    hazards: Iterable[tuple[int, int]],
    hazard_count: dict[tuple[int, int], int],
    hazard_damage: int,
    deadline: float | None,
  ) -> None:
    self.board = board
    self.deadline = deadline
    self.hazard_damage = hazard_damage
    self.food_bits = board.set_to_bits(set(food))
    self.hazard_bits = board.set_to_bits(set(hazards))
    self.hazard_stacks = {
      board.idx(x, y): count for (x, y), count in hazard_count.items()
    }
    self.transposition: dict[tuple[DuelState, int], tuple[float, str]] = {}
    self.killer_moves: dict[int, int] = {}
    self.history: dict[int, int] = {}
    self.nodes = 0
    self.cache_hits = 0

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

  def search(
    self,
    my_body: list[dict],
    enemy_body: list[dict],
    my_health: int,
    enemy_health: int,
    max_depth: int,
    previous_hazards: Iterable[tuple[int, int]],
  ) -> tuple[float, int]:
    state = DuelState(
      my_body=self.body_from_dicts(my_body),
      enemy_body=self.body_from_dicts(enemy_body),
      food_bits=self.food_bits,
      my_health=my_health,
      enemy_health=enemy_health,
      previous_hazard_bits=self.board.set_to_bits(set(previous_hazards)),
    )
    result = self._evaluate(state)
    completed_depth = 0

    for depth in range(1, max_depth + 1):
      if self._out_of_time(5.0):
        break
      value, completed = self._search(state, depth, -float("inf"), float("inf"))
      if not completed:
        break
      result = value
      completed_depth = depth

    return result, completed_depth

  def search_depth(
    self,
    my_body: list[dict],
    enemy_body: list[dict],
    my_health: int,
    enemy_health: int,
    depth: int,
    previous_hazards: Iterable[tuple[int, int]],
  ) -> float:
    state = DuelState(
      my_body=self.body_from_dicts(my_body),
      enemy_body=self.body_from_dicts(enemy_body),
      food_bits=self.food_bits,
      my_health=my_health,
      enemy_health=enemy_health,
      previous_hazard_bits=self.board.set_to_bits(set(previous_hazards)),
    )
    value, _ = self._search(state, depth, -float("inf"), float("inf"))
    return value

  def _search(self, state: DuelState, depth: int, alpha: float, beta: float) -> tuple[float, bool]:
    self.nodes += 1
    if self._out_of_time():
      return self._evaluate(state), False
    if depth <= 0:
      return self._evaluate(state), True

    cache_key = (state, depth)
    original_alpha, original_beta = alpha, beta
    cached = self.transposition.get(cache_key)
    if cached is not None:
      self.cache_hits += 1
      cached_value, bound = cached
      if bound == "exact":
        return cached_value, True
      if bound == "lower":
        alpha = max(alpha, cached_value)
      else:
        beta = min(beta, cached_value)
      if alpha >= beta:
        return cached_value, True

    my_moves = self._legal_targets(state.my_body, state.enemy_body)
    enemy_moves = self._legal_targets(state.enemy_body, state.my_body)
    if not my_moves:
      return self.LOSS - depth, True
    if not enemy_moves:
      return self.WIN + depth, True

    my_moves = self._ordered_moves(my_moves, state, depth, True)
    enemy_moves = self._ordered_moves(enemy_moves, state, depth, False)
    best = -float("inf")

    for my_target in my_moves:
      worst = float("inf")
      for enemy_target in enemy_moves:
        if self._out_of_time():
          return (best if best != -float("inf") else self._evaluate(state)), False
        child, terminal = self._advance(state, my_target, enemy_target)
        if terminal is not None:
          value = terminal
          completed = True
        else:
          value, completed = self._search(child, depth - 1, alpha, beta)
        if not completed:
          return (best if best != -float("inf") else value), False
        worst = min(worst, value)
        if worst <= alpha:
          self.killer_moves[depth] = my_target
          self.history[my_target] = self.history.get(my_target, 0) + depth * depth
          break

      best = max(best, worst)
      alpha = max(alpha, best)
      if alpha >= beta:
        break

    if best <= original_alpha:
      bound = "upper"
    elif best >= original_beta:
      bound = "lower"
    else:
      bound = "exact"
    self.transposition[cache_key] = (best, bound)
    return best, True

  def _advance(self, state: DuelState, my_target: int, enemy_target: int) -> tuple[DuelState, float | None]:
    my_ate = bool((1 << my_target) & state.food_bits)
    enemy_ate = bool((1 << enemy_target) & state.food_bits)
    my_body = self._advance_body(state.my_body, my_target, my_ate)
    enemy_body = self._advance_body(state.enemy_body, enemy_target, enemy_ate)

    my_dead = my_target in my_body[1:] or my_target in enemy_body[1:]
    enemy_dead = enemy_target in enemy_body[1:] or enemy_target in my_body[1:]

    if my_target == enemy_target:
      if len(my_body) <= len(enemy_body):
        my_dead = True
      if len(enemy_body) <= len(my_body):
        enemy_dead = True

    my_health = 100 if my_ate else state.my_health - 1
    enemy_health = 100 if enemy_ate else state.enemy_health - 1
    if not my_ate:
      my_health -= self._hazard_cost(my_target, state.previous_hazard_bits)
    if not enemy_ate:
      enemy_health -= self._hazard_cost(enemy_target, state.previous_hazard_bits)
    my_dead = my_dead or my_health <= 0
    enemy_dead = enemy_dead or enemy_health <= 0

    if my_dead and enemy_dead:
      return state, -500.0
    if my_dead:
      return state, self.LOSS
    if enemy_dead:
      return state, self.WIN

    eaten_bits = 0
    if my_ate:
      eaten_bits |= 1 << my_target
    if enemy_ate:
      eaten_bits |= 1 << enemy_target
    child = DuelState(
      my_body=my_body,
      enemy_body=enemy_body,
      food_bits=state.food_bits & ~eaten_bits,
      my_health=my_health,
      enemy_health=enemy_health,
      previous_hazard_bits=self.hazard_bits,
    )
    return child, None

  def _legal_targets(self, body: Body, other_body: Body) -> list[int]:
    occupied = self._body_bits(body) | self._body_bits(other_body)
    if not self._tail_stacked(body):
      occupied &= ~(1 << body[-1])
    if not self._tail_stacked(other_body):
      occupied &= ~(1 << other_body[-1])
    legal = self.board.neighbors_of(body[0]) & ~occupied & self.board.board_mask
    return list(self._iter_bits(legal))

  def _ordered_moves(self, moves: list[int], state: DuelState, depth: int, mine: bool) -> list[int]:
    body = state.my_body if mine else state.enemy_body
    other = state.enemy_body if mine else state.my_body
    killer = self.killer_moves.get(depth)
    center_x = (self.board.width - 1) / 2.0
    center_y = (self.board.height - 1) / 2.0

    def score(target: int) -> tuple[float, int]:
      x, y = self.board.coord(target)
      food_bonus = 200.0 if (1 << target) & state.food_bits else 0.0
      space = self.board.flood_count(target, (self._body_bits(body[1:]) | self._body_bits(other[1:])) & ~(1 << target))
      center = -(abs(x - center_x) + abs(y - center_y))
      killer_bonus = 10_000.0 if target == killer else 0.0
      return killer_bonus + self.history.get(target, 0) + food_bonus + space * 2.0 + center, -target

    # Our strongest-looking moves first; enemy ordering uses the same quality
    # estimate because dangerous enemy replies tend to gain space and food.
    return sorted(moves, key=score, reverse=True)

  def _evaluate(self, state: DuelState) -> float:
    my_blocked = self._body_bits(state.my_body[1:]) | self._body_bits(state.enemy_body[1:])
    my_space = self.board.flood_count(state.my_body[0], my_blocked)
    enemy_space = self.board.flood_count(state.enemy_body[0], my_blocked)
    my_liberties = self.board.open_neighbor_count(state.my_body[0], my_blocked)
    enemy_liberties = self.board.open_neighbor_count(state.enemy_body[0], my_blocked)
    length_score = (len(state.my_body) - len(state.enemy_body)) * 18.0
    health_score = (state.my_health - state.enemy_health) * 0.15
    return (my_space - enemy_space) * 2.0 + (my_liberties - enemy_liberties) * 12.0 + length_score + health_score

  def _hazard_cost(self, target: int, previous_hazard_bits: int) -> int:
    bit = 1 << target
    if not (bit & self.hazard_bits & previous_hazard_bits):
      return 0
    return self.hazard_damage * self.hazard_stacks.get(target, 1)

  @staticmethod
  def _advance_body(body: Body, target: int, ate: bool) -> Body:
    return (target,) + body if ate else (target,) + body[:-1]

  @staticmethod
  def _tail_stacked(body: Body) -> bool:
    return len(body) >= 2 and body[-1] == body[-2]

  @staticmethod
  def _body_bits(body: Body) -> int:
    bits = 0
    for cell in body:
      bits |= 1 << cell
    return bits

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

  def _out_of_time(self, reserve_ms: float = 0.0) -> bool:
    if self.deadline is None:
      return False
    return perf_counter() + reserve_ms / 1000.0 >= self.deadline