"""Reusable bitboard and adversarial survival rollouts."""

from __future__ import annotations

from time import perf_counter

from snakes.engine.survival_search import CompactSurvivalSearch

class BitboardSurvivalMixin:
  def _future_rollout_bonus(
      self, move: str, safe_moves: dict, my_body: list, other_snakes: list,
      food_set: set, is_constrictor: bool, width: int, height: int,
      enemy_can_grow: dict, deadline: float | None,
    ) -> float:
      pos = safe_moves.get(move)
      if pos is None:
        return -250.0
      # Duel minimax already advances the opponent exactly. Keep the much faster
      # bitboard-native solo rollout here instead of paying for the same response
      # model twice. Constrictor and multiplayer still use adversarial rollouts.
      if len(other_snakes) == 1 and not is_constrictor:
        return super()._future_rollout_bonus(
          move, safe_moves, my_body, other_snakes, food_set, is_constrictor,
          width, height, enemy_can_grow, deadline,
        )
      if self._survival_search_context is None:
        remaining = self._remaining_ms(deadline)
        enemy_branch = 2 if len(other_snakes) <= 2 and remaining > 100 else 1
        self._survival_search_context = CompactSurvivalSearch(
          board=self._get_bb(width, height),
          food=food_set,
          is_constrictor=is_constrictor,
          deadline=deadline,
          branch=self._planning_branch,
          enemy_branch=enemy_branch,
          response_cap=8 if remaining > 150 else 4,
        )
      remaining = self._remaining_ms(deadline)
      depth = min(self._planning_depth, 2 if len(other_snakes) > 1 else 3)
      if remaining < 90:
        depth = min(depth, 2)
      elif remaining > 250 and len(other_snakes) <= 2:
        depth = min(4, depth + 1)
      raw = self._survival_search_context.search_selected(
        my_body=my_body,
        enemies=other_snakes,
        target=(pos["x"], pos["y"]),
        depth=depth,
      )
      return raw * 0.15

  def _future_position_score(
      self, my_body: list, other_snakes: list, food_set: set, is_constrictor: bool,
      width: int, height: int, enemy_can_grow: dict, deadline: float | None,
    ) -> float:
      """S9: Bitboard-native position scoring for the survival tree.

      Builds blocked bitboard directly from body lists (no intermediate set).
      Uses precomputed enemy bits instead of rebuilding attack map per node.
      """
      if deadline is not None and perf_counter() >= deadline:
        return 0.0

      bb = self._bb  # already initialised in choose_move
      w = bb.width
      head = my_body[0]
      hx, hy = head["x"], head["y"]
      head_idx = hy * w + hx
      head_bit = 1 << head_idx
      body_len = len(my_body)

      # ── Build blocked bitboard directly (no set) ──────────────────────
      my_bits = 0
      for seg in my_body:
        my_bits |= 1 << (seg["y"] * w + seg["x"])

      # Own tail vacates unless stacked or constrictor
      if not is_constrictor and body_len >= 2:
        t, t2 = my_body[-1], my_body[-2]
        if not (t["x"] == t2["x"] and t["y"] == t2["y"]):
          my_bits &= ~(1 << (t["y"] * w + t["x"]))

      # Enemy body (precomputed) minus vacating tails
      en_bits = self._enemy_body_bits & ~self._enemy_tail_bits

      blocked_bits = (my_bits | en_bits) & ~head_bit

      # ── Reachable space ───────────────────────────────────────────────
      reachable = bb.flood_count(head_idx, blocked_bits)
      required = body_len + max(3, body_len // 6) if is_constrictor else body_len
      if reachable < required:
        return -5000.0

      # ── Open neighbours (liberties) ───────────────────────────────────
      nb_free = bb._neighbor_masks[head_idx] & ~blocked_bits & bb.board_mask
      liberties = nb_free.bit_count()
      if liberties == 0:
        return -5000.0

      # ── Safe next options (enemy-attack aware) ────────────────────────
      # Rebuild danger for the simulated length. The root-turn danger mask is
      # stale after eating and includes enemy moves blocked in this future body.
      danger_here = 0
      for enemy in other_snakes:
        enemy_len = enemy.get("length", len(enemy["body"]))
        if enemy_len < body_len:
          continue
        enemy_head = enemy["head"]
        enemy_idx = enemy_head["y"] * w + enemy_head["x"]
        danger_here |= bb._neighbor_masks[enemy_idx]
      danger_here &= ~blocked_bits
      safe_nb = nb_free & ~danger_here
      en_safe = safe_nb.bit_count()

      if en_safe == 0:
        return -4000.0

      next_opts = liberties
      sc = reachable * 1.9 + liberties * 14.0 + next_opts * 11.0 + en_safe * 26.0
      if en_safe == 1:
        sc -= 420.0
      return sc

  def _future_survival_tree(
      self, my_body: list, other_snakes: list, food_set: set, is_constrictor: bool,
      width: int, height: int, enemy_can_grow: dict,
      depth: int, branch: int, deadline: float | None,
    ) -> float:
      """S9/S11: Bitboard-accelerated survival tree.

      Inlines legal-move check with bitboard ops instead of per-direction
      Python loops.  Uses the bitboard-native _future_position_score.
      """
      if depth <= 0 or (deadline is not None and perf_counter() >= deadline):
        return 0.0

      bb = self._bb
      w = bb.width
      head = my_body[0]
      hx, hy = head["x"], head["y"]
      head_idx = hy * w + hx
      body_len = len(my_body)

      # ── Build occupied bitboard for legal-move check ──────────────────
      occupied_bits = 0
      for seg in my_body:
        occupied_bits |= 1 << (seg["y"] * w + seg["x"])
      occupied_bits |= self._enemy_body_bits

      # Own tail can be stepped on if not stacked/constrictor
      passable = 0
      if not is_constrictor and body_len >= 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 vacating tails are also steppable
      passable |= self._enemy_tail_bits

      # Legal moves: free neighbours OR passable tiles
      legal_bits = bb._neighbor_masks[head_idx] & ((~occupied_bits & bb.board_mask) | passable)

      if not legal_bits:
        return -5000.0

      # ── Precompute food bitboard once ─────────────────────────────────
      food_bits_local = 0
      for fx, fy in food_set:
        food_bits_local |= 1 << (fy * w + fx)

      # ── Score each legal move ─────────────────────────────────────────
      scored: list[tuple[float, list]] = []
      temp = legal_bits
      while temp:
        if deadline is not None and perf_counter() >= deadline:
          break
        bit = temp & (-temp)
        temp ^= bit
        idx = bit.bit_length() - 1
        nx, ny = idx % w, idx // w
        pos = {"x": nx, "y": ny}
        ate = bool(bit & food_bits_local)
        fb = self._future_body(my_body, pos, ate, is_constrictor)
        sc = self._future_position_score(
          fb, other_snakes, food_set, is_constrictor,
          width, height, enemy_can_grow, deadline,
        )
        scored.append((sc, fb))

      if not scored:
        return -5000.0

      DEATH = self._TREE_DEATH_THRESHOLD
      viable = [(sc, fb) for sc, fb in scored if sc > DEATH]
      if not viable:
        return max(sc for sc, _ in scored)

      viable.sort(key=lambda x: x[0], reverse=True)

      if depth == 1:
        return viable[0][0]

      best = viable[0][0]
      for sc, fb in viable[:branch]:
        if deadline is not None and perf_counter() >= deadline:
          break
        cont = self._future_survival_tree(
          fb, other_snakes, food_set, is_constrictor,
          width, height, enemy_can_grow, depth - 1, branch, deadline,
        )
        total = sc + cont * 0.72
        if total > best:
          best = total
      return best