"""Bitboard-backed spatial primitives shared by competitive snakes."""

from __future__ import annotations

from snakes.engine.bitboard import BitBoard

class BitboardSpatialMixin:
  def _get_bb(self, width: int, height: int) -> BitBoard:
      """Return (possibly cached) BitBoard for the current dimensions."""
      if self._bb is None or width != self._bb_w or height != self._bb_h:
        self._bb = BitBoard(width, height)
        self._bb_w = width
        self._bb_h = height
      return self._bb

  def _blocked_to_bits(self, blocked: set[tuple[int, int]], width: int, height: int) -> int:
      """Convert blocked cells to bits without stale identity-based caching."""
      return self._get_bb(width, height).set_to_bits(blocked)

  def _flood_fill_count(self, start: tuple, blocked: set, width: int, height: int) -> int:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      start_idx = bb.idx(start[0], start[1])

      # A7/E2: per-turn transposition cache (kept from Apex)
      cache_key = (start_idx, blocked_bits, width, height)
      cached = self._bfs_cache.get(cache_key)
      if cached is not None:
        return cached

      result = bb.flood_count(start_idx, blocked_bits)

      if len(self._bfs_cache) < self._bfs_cache_max:
        self._bfs_cache[cache_key] = result
      return result

  def _territory_fast(
      self, my_pos: tuple, blocked: set, width: int, height: int,
      deadline: float | None = None,
    ) -> int:
      if not self._enemy_heads:
        return 0
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      my_idx = bb.idx(my_pos[0], my_pos[1])
      enemy_idxs = [bb.idx(eh[0], eh[1]) for eh in self._enemy_heads]
      return bb.territory(my_idx, enemy_idxs, blocked_bits)

  def _articulation_penalty(
      self, point: tuple, blocked: set, width: int, height: int, required_space: int,
    ) -> float:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      point_idx = bb.idx(point[0], point[1])

      sizes = bb.partition_sizes(point_idx, blocked_bits)
      if not sizes:
        return 0.0

      min_size = min(sizes)
      if min_size < required_space:
        return 1500.0
      elif min_size < required_space * 2:
        return 400.0
      else:
        return 85.0

  def _bounded_bfs(self, start: tuple, blocked: set, width: int, height: int, limit: int) -> set:
      """Bitboard-accelerated bounded BFS. Returns a set for API compatibility."""
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      start_idx = bb.idx(start[0], start[1])
      reachable_bits = bb.flood_fill(start_idx, blocked_bits)

      result: set[tuple[int, int]] = set()
      temp = reachable_bits
      w = bb.width
      while temp:
        bit = temp & (-temp)
        idx = bit.bit_length() - 1
        result.add((idx % w, idx // w))
        temp ^= bit
        if len(result) >= limit:
          break
      return result

  def _distance_map(self, start: tuple, blocked: set, width: int, height: int) -> dict:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      start_idx = bb.idx(start[0], start[1])
      idx_dmap = bb.distance_map(start_idx, blocked_bits)
      w = bb.width
      return {(idx % w, idx // w): d for idx, d in idx_dmap.items()}

  def _path_distance(
      self, start: tuple, goal: tuple, blocked: set, width: int, height: int,
    ) -> int | None:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      return bb.path_distance(
        bb.idx(start[0], start[1]),
        bb.idx(goal[0], goal[1]),
        blocked_bits,
      )

  def _nearest_food_info(
      self, start: tuple, food_set: set, blocked: set, width: int, height: int,
    ) -> tuple[int | None, tuple | None]:
      if not food_set:
        return None, None
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      food_bits = bb.set_to_bits(food_set)
      start_idx = bb.idx(start[0], start[1])
      dist, cell_idx = bb.nearest_food(start_idx, food_bits, blocked_bits)
      if dist is None or cell_idx is None:
        return None, None
      return dist, bb.coord(cell_idx)

  def _open_neighbor_count(self, start: tuple, blocked: set, width: int, height: int) -> int:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      return bb.open_neighbor_count(bb.idx(start[0], start[1]), blocked_bits)

  def _next_turn_options(self, head: dict, blocked: set, width: int, height: int) -> int:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      return bb.open_neighbor_count(bb.idx(head["x"], head["y"]), blocked_bits)

  def _legal_moves(
      self, my_head, my_body: list, other_snakes: list,
      food_set: set, is_constrictor: bool, width: int, height: int,
      enemy_can_grow: dict | None = None,
    ):
      """S10: Bitboard-accelerated legal move generation."""
      bb = self._get_bb(width, height)
      w = bb.width

      # Build occupied bitboard
      occupied = 0
      for seg in my_body:
        occupied |= 1 << (seg["y"] * w + seg["x"])
      for snake in other_snakes:
        for seg in snake["body"]:
          occupied |= 1 << (seg["y"] * w + seg["x"])

      hx, hy = my_head["x"], my_head["y"]
      head_idx = hy * w + hx

      # Own tail can be stepped on
      passable = 0
      if not is_constrictor and len(my_body) >= 2:
        t, t2 = my_body[-1], my_body[-2]
        if not (t["x"] == t2["x"] and t["y"] == t2["y"]):
          passable |= 1 << (t["y"] * w + t["x"])

      # Enemy tails that will vacate
      if not is_constrictor:
        for snake in other_snakes:
          sbody = snake["body"]
          if len(sbody) < 2:
            continue
          st, st2 = sbody[-1], sbody[-2]
          if st["x"] == st2["x"] and st["y"] == st2["y"]:
            continue  # stacked
          sid = snake.get("id")
          can_grow = None
          if enemy_can_grow is not None and sid is not None:
            can_grow = enemy_can_grow.get(sid)
          if can_grow is None:
            can_grow = self._enemy_can_grow_this_turn(snake, food_set)
          if not can_grow:
            passable |= 1 << (st["y"] * w + st["x"])

      legal = bb._neighbor_masks[head_idx] & ((~occupied & bb.board_mask) | passable)

      safe: dict[str, dict[str, int]] = {}
      for name, (dx, dy) in self.DIRECTIONS.items():
        nx, ny = hx + dx, hy + dy
        if 0 <= nx < w and 0 <= ny < bb.height:
          if (1 << (ny * w + nx)) & legal:
            safe[name] = {"x": nx, "y": ny}
      return safe

  def _enemy_confinement_metrics(
      self, enemy_head: tuple, blocked: set, width: int, height: int,
    ) -> tuple[int, int]:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      eh_idx = bb.idx(enemy_head[0], enemy_head[1])
      eb_bits = blocked_bits & ~(1 << eh_idx)
      space = bb.flood_count(eh_idx, eb_bits)
      options = bb.open_neighbor_count(eh_idx, eb_bits)
      return space, options

  def _enemy_constrictor_projection(
      self, other_snakes: list, blocked: set, width: int, height: int,
    ) -> tuple[int, int]:
      bb = self._get_bb(width, height)
      blocked_bits = self._blocked_to_bits(blocked, width, height)
      best_space = 0
      total_opts = 0
      for enemy in other_snakes:
        eh = (enemy["head"]["x"], enemy["head"]["y"])
        eh_idx = bb.idx(eh[0], eh[1])
        nb = bb.neighbors_of(eh_idx) & ~blocked_bits & bb.board_mask
        temp = nb
        while temp:
          total_opts += 1
          bit = temp & (-temp)
          n_idx = bit.bit_length() - 1
          sp = bb.flood_count(n_idx, blocked_bits | bit)
          if sp > best_space:
            best_space = sp
          temp ^= bit
      return best_space, total_opts