from collections import deque
from typing import Any, cast
import random
import os

from snakes.TemplateSnake import TemplateSnake

class BestBattleSnake(TemplateSnake):
  DIRECTIONS = {
    "up": (0, 1),
    "down": (0, -1),
    "left": (-1, 0),
    "right": (1, 0),
  }

  OPPOSITE = {
    "up": "down",
    "down": "up",
    "left": "right",
    "right": "left",
  }

  def __init__(self):
    super().__init__()
    self.name = "BestBattleSnake"
    self.recent_heads = deque(maxlen=14)
    self.last_move = None
    self.last_game_id = None
    self.duel_style = self._get_duel_style()

  def _get_duel_style(self):
    value = os.getenv("BATTLE_SNAKE_DUEL_STYLE")
    if value is None:
      value = os.getenv("DUEL_STYLE", "balanced")

    style = value.strip().lower()
    if style not in {"safe", "balanced", "aggressive"}:
      return "balanced"
    return style

  def _duel_weights(self, style):
    if style == "safe":
      return {
        "head_pressure": 0.65,
        "distance_safety": 1.30,
        "food_bias": 1.00,
      }
    if style == "aggressive":
      return {
        "head_pressure": 1.35,
        "distance_safety": 0.75,
        "food_bias": 0.85,
      }
    return {
      "head_pressure": 1.00,
      "distance_safety": 1.00,
      "food_bias": 1.00,
    }

  def choose_move(self, game_data):
    self.game_board = game_data
    self.calculations = []
    self.duel_style = self._get_duel_style()

    game_id = getattr(game_data, "id", None)
    turn = game_data.get_turn()
    if game_id != self.last_game_id or turn <= 1:
      self.recent_heads.clear()
      self.last_move = None
      self.last_game_id = game_id

    my_snake = cast(dict[str, Any], game_data.get_my_snake())
    my_head = my_snake["head"]
    my_body = my_snake["body"]
    my_len = my_snake.get("length", len(my_body))
    my_health = my_snake.get("health", 100)

    width = game_data.get_width()
    height = game_data.get_height()
    board_area = max(1, width * height)
    occupancy_ratio = my_len / board_area
    preserve_space_mode = occupancy_ratio >= 0.34 and my_health > 35
    foods = game_data.get_food()
    hazards = game_data.get_hazard()
    other_snakes = game_data.get_other_snakes()
    is_constrictor = game_data.get_type() == "constrictor"

    food_set = {(food["x"], food["y"]) for food in foods}
    hazard_set = {(hazard["x"], hazard["y"]) for hazard in hazards}
    current_head_point = (my_head["x"], my_head["y"])

    safe_moves = self._legal_moves(
      my_head=my_head,
      my_body=my_body,
      other_snakes=other_snakes,
      food_set=food_set,
      is_constrictor=is_constrictor,
      width=width,
      height=height,
    )

    if not safe_moves:
      fallback = self._fallback_move(my_head, width, height)
      self.recent_heads.append(current_head_point)
      self.last_move = fallback
      self.add_to_history(
        {
          "turn": turn,
          "move": fallback,
          "reason": "no_safe_moves",
        }
      )
      return fallback

    enemy_attack_map = self._build_enemy_attack_map(
      my_snake=my_snake,
      other_snakes=other_snakes,
      food_set=food_set,
      is_constrictor=is_constrictor,
      width=width,
      height=height,
    )

    if is_constrictor:
      best_move, scores = self._choose_constrictor_move(
        safe_moves=safe_moves,
        my_body=my_body,
        my_len=my_len,
        other_snakes=other_snakes,
        food_set=food_set,
        enemy_attack_map=enemy_attack_map,
        width=width,
        height=height,
      )
      self.recent_heads.append(current_head_point)
      self.last_move = best_move
      self.add_to_history({"turn": turn, "move": best_move, "scores": scores})
      return best_move

    if len(other_snakes) == 1:
      best_move, scores = self._choose_duel_move(
        safe_moves=safe_moves,
        my_body=my_body,
        my_len=my_len,
        my_health=my_health,
        foods=foods,
        food_set=food_set,
        hazards=hazards,
        hazard_set=hazard_set,
        other_snakes=other_snakes,
        enemy_attack_map=enemy_attack_map,
        width=width,
        height=height,
      )
      self.recent_heads.append(current_head_point)
      self.last_move = best_move
      self.add_to_history({"turn": turn, "move": best_move, "scores": scores})
      return best_move

    scores: dict[str, float] = {}
    move_safety: dict[str, dict[str, Any]] = {}
    for move, pos in safe_moves.items():
      point = (pos["x"], pos["y"])
      ate_food = point in food_set

      future_body = self._future_body(my_body, pos, ate_food, is_constrictor)
      blocked = self._simulation_blocked(
        future_body=future_body,
        other_snakes=other_snakes,
        food_set=food_set,
        is_constrictor=is_constrictor,
      )
      blocked.discard(point)

      reachable_space = self._flood_fill_count(point, blocked, width, height)
      required_space = len(future_body)
      liberties = self._open_neighbor_count(point, blocked, width, height)
      next_options = self._next_turn_option_count(
        future_body, blocked, width, height
      )
      territory = self._territory_control_score(
        my_start=point,
        enemy_starts=[
          (snake["head"]["x"], snake["head"]["y"]) for snake in other_snakes
        ],
        blocked=blocked,
        width=width,
        height=height,
      )

      nearest_food_dist = self._nearest_food_distance(
        point, foods, blocked, width, height
      )
      future_tail = future_body[-1]
      tail_point = (future_tail["x"], future_tail["y"])
      tail_dist = self._path_distance(
        point, tail_point, blocked - {tail_point}, width, height
      )
      has_tail_escape = tail_dist is not None
      likely_dead_end = (
        (reachable_space < required_space and not has_tail_escape)
        or (liberties == 0 and not has_tail_escape)
        or (next_options == 0 and not has_tail_escape)
      )

      score = 0.0
      score += reachable_space * 2.6
      score += liberties * 18.0
      score += next_options * 10.0
      score += territory * 0.35

      if reachable_space < required_space:
        score -= 1200.0
      if liberties == 0:
        score -= 900.0

      enemy_len = enemy_attack_map.get(point)
      if enemy_len is not None:
        if enemy_len >= my_len:
          score -= 1200.0
        else:
          score += 70.0

      hunger_urgency = max(0.0, (60.0 - my_health) / 60.0)
      if nearest_food_dist is not None:
        score += (28.0 + 70.0 * hunger_urgency) / (nearest_food_dist + 1)
      elif my_health < 30:
        score -= 150.0

      if ate_food:
        if likely_dead_end:
          score -= 1800.0
        else:
          score += 260.0 + 220.0 * hunger_urgency

      if preserve_space_mode and ate_food and my_health > 45:
        score -= 280.0

      if tail_dist is not None:
        score += 12.0 / (tail_dist + 1)
      else:
        score -= 40.0

      if point in hazard_set:
        score -= 70.0 if my_health > 35 else 250.0

      score -= self._revisit_penalty(point)

      if self.last_move == move:
        score += 6.0
      elif (
        self.last_move
        and self.OPPOSITE[self.last_move] == move
        and len(safe_moves) > 1
      ):
        score -= 20.0

      health_after_move = 100 if ate_food else my_health - 1
      if point in hazard_set:
        health_after_move -= 15
      if health_after_move <= 0:
        score -= 10000.0

      is_losing_head_to_head = enemy_len is not None and enemy_len >= my_len
      is_dead_end = likely_dead_end
      move_safety[move] = {
        "is_survivable": (not is_dead_end)
        and (not is_losing_head_to_head)
        and health_after_move > 0,
        "reachable_space": reachable_space,
        "next_options": next_options,
        "tail_escape": has_tail_escape,
      }

      scores[move] = round(score, 5)

    survivable_moves = [
      move for move, data in move_safety.items() if data["is_survivable"]
    ]
    if survivable_moves:
      best_space = max(
        move_safety[move]["reachable_space"] for move in survivable_moves
      )
      roomy_moves = [
        move
        for move in survivable_moves
        if move_safety[move]["reachable_space"]
        >= max(1, int(best_space * 0.60))
      ]
      tail_escape_moves = [
        move for move in survivable_moves if move_safety[move]["tail_escape"]
      ]
      if tail_escape_moves:
        considered_moves = tail_escape_moves
      else:
        considered_moves = roomy_moves if roomy_moves else survivable_moves
    else:
      considered_moves = list(scores.keys())

    best_score = max(scores[move] for move in considered_moves)
    top_moves = [
      move for move in considered_moves if best_score - scores[move] <= 1.5
    ]
    best_move = random.choice(top_moves)
    self.recent_heads.append(current_head_point)
    self.last_move = best_move
    self.add_to_history({"turn": turn, "move": best_move, "scores": scores})
    return best_move

  def _choose_duel_move(
    self,
    safe_moves,
    my_body,
    my_len,
    my_health,
    foods,
    food_set,
    hazards,
    hazard_set,
    other_snakes,
    enemy_attack_map,
    width,
    height,
  ):
    duel_weights = self._duel_weights(self.duel_style)
    enemy = other_snakes[0]
    enemy_head = (enemy["head"]["x"], enemy["head"]["y"])
    enemy_len = enemy.get("length", len(enemy["body"]))

    scores: dict[str, float] = {}
    move_safety: dict[str, dict[str, Any]] = {}

    for move, pos in safe_moves.items():
      point = (pos["x"], pos["y"])
      ate_food = point in food_set

      future_body = self._future_body(
        my_body, pos, ate_food, is_constrictor=False
      )
      blocked = self._simulation_blocked(
        future_body=future_body,
        other_snakes=other_snakes,
        food_set=food_set,
        is_constrictor=False,
      )
      blocked.discard(point)

      reachable_space = self._flood_fill_count(point, blocked, width, height)
      required_space = len(future_body)
      liberties = self._open_neighbor_count(point, blocked, width, height)
      next_options = self._next_turn_option_count(
        future_body, blocked, width, height
      )

      nearest_food_dist = self._nearest_food_distance(
        point, foods, blocked, width, height
      )
      future_tail = future_body[-1]
      tail_point = (future_tail["x"], future_tail["y"])
      tail_dist = self._path_distance(
        point, tail_point, blocked - {tail_point}, width, height
      )
      territory = self._territory_control_score(
        my_start=point,
        enemy_starts=[enemy_head],
        blocked=blocked,
        width=width,
        height=height,
      )

      has_tail_escape = tail_dist is not None
      likely_dead_end = (
        (reachable_space < required_space and not has_tail_escape)
        or (liberties == 0 and not has_tail_escape)
        or (next_options == 0 and not has_tail_escape)
      )

      enemy_attack_len = enemy_attack_map.get(point)
      losing_head_to_head = (
        enemy_attack_len is not None and enemy_attack_len >= my_len
      )
      direct_head_distance = self._manhattan(point, enemy_head)

      score = 0.0
      score += reachable_space * 2.8
      score += liberties * 18.0
      score += next_options * 10.0
      score += territory * 0.50

      if likely_dead_end:
        score -= 1400.0

      if losing_head_to_head:
        score -= 1500.0

      if my_len > enemy_len:
        if direct_head_distance == 1:
          score += 220.0 * duel_weights["head_pressure"]
        elif direct_head_distance == 2:
          score += 80.0 * duel_weights["head_pressure"]
      else:
        if direct_head_distance <= 2:
          score -= 120.0 * duel_weights["distance_safety"]

      hunger_urgency = max(0.0, (65.0 - my_health) / 65.0)
      if nearest_food_dist is not None:
        score += (
          (25.0 + 90.0 * hunger_urgency) * duel_weights["food_bias"]
        ) / (nearest_food_dist + 1)

      if ate_food:
        if likely_dead_end:
          score -= 1700.0
        else:
          score += 260.0 + 250.0 * hunger_urgency

      if tail_dist is not None:
        score += 14.0 / (tail_dist + 1)
      else:
        score -= 50.0

      if point in hazard_set:
        score -= 70.0 if my_health > 35 else 250.0

      score -= self._revisit_penalty(point)

      if self.last_move == move:
        score += 6.0
      elif (
        self.last_move
        and self.OPPOSITE[self.last_move] == move
        and len(safe_moves) > 1
      ):
        score -= 20.0

      health_after_move = 100 if ate_food else my_health - 1
      if point in hazard_set:
        health_after_move -= 15
      if health_after_move <= 0:
        score -= 10000.0

      move_safety[move] = {
        "is_survivable": (not likely_dead_end)
        and (not losing_head_to_head)
        and health_after_move > 0,
        "reachable_space": reachable_space,
        "tail_escape": has_tail_escape,
      }
      scores[move] = round(score, 5)

    survivable_moves = [
      move for move, data in move_safety.items() if data["is_survivable"]
    ]
    if survivable_moves:
      tail_escape_moves = [
        move for move in survivable_moves if move_safety[move]["tail_escape"]
      ]
      if tail_escape_moves:
        considered_moves = tail_escape_moves
      else:
        best_space = max(
          move_safety[move]["reachable_space"] for move in survivable_moves
        )
        considered_moves = [
          move
          for move in survivable_moves
          if move_safety[move]["reachable_space"]
          >= max(1, int(best_space * 0.60))
        ]
        if not considered_moves:
          considered_moves = survivable_moves
    else:
      considered_moves = list(scores.keys())

    best_score = max(scores[move] for move in considered_moves)
    top_moves = [
      move for move in considered_moves if best_score - scores[move] <= 1.5
    ]
    return random.choice(top_moves), scores

  def _choose_constrictor_move(
    self,
    safe_moves,
    my_body,
    my_len,
    other_snakes,
    food_set,
    enemy_attack_map,
    width,
    height,
  ):
    scores: dict[str, float] = {}
    move_safety: dict[str, dict[str, Any]] = {}

    for move, pos in safe_moves.items():
      point = (pos["x"], pos["y"])
      future_body = self._future_body(
        my_body, pos, ate_food=False, is_constrictor=True
      )
      blocked = self._simulation_blocked(
        future_body=future_body,
        other_snakes=other_snakes,
        food_set=food_set,
        is_constrictor=True,
      )
      blocked.discard(point)

      reachable_space = self._flood_fill_count(point, blocked, width, height)
      required_space = len(future_body) + 1
      liberties = self._open_neighbor_count(point, blocked, width, height)
      next_options = self._next_turn_option_count(
        future_body, blocked, width, height
      )
      territory = self._territory_control_score(
        my_start=point,
        enemy_starts=[
          (snake["head"]["x"], snake["head"]["y"]) for snake in other_snakes
        ],
        blocked=blocked,
        width=width,
        height=height,
      )

      enemy_len = enemy_attack_map.get(point)
      is_losing_head_to_head = enemy_len is not None and enemy_len >= my_len

      is_dead_end = (
        reachable_space < required_space or liberties == 0 or next_options == 0
      )

      score = 0.0
      score += reachable_space * 3.8
      score += liberties * 24.0
      score += next_options * 16.0
      score += territory * 0.65

      if is_dead_end:
        score -= 2600.0

      if is_losing_head_to_head:
        score -= 2400.0
      elif enemy_len is not None:
        score += 90.0

      score -= self._revisit_penalty(point)

      if self.last_move == move:
        score += 10.0
      elif (
        self.last_move
        and self.OPPOSITE[self.last_move] == move
        and len(safe_moves) > 1
      ):
        score -= 35.0

      move_safety[move] = {
        "is_survivable": (not is_dead_end) and (not is_losing_head_to_head),
        "reachable_space": reachable_space,
      }
      scores[move] = round(score, 5)

    survivable_moves = [
      move for move, data in move_safety.items() if data["is_survivable"]
    ]
    if survivable_moves:
      best_space = max(
        move_safety[move]["reachable_space"] for move in survivable_moves
      )
      considered_moves = [
        move
        for move in survivable_moves
        if move_safety[move]["reachable_space"]
        >= max(1, int(best_space * 0.70))
      ]
      if not considered_moves:
        considered_moves = survivable_moves
    else:
      considered_moves = list(scores.keys())

    best_score = max(scores[move] for move in considered_moves)
    top_moves = [
      move for move in considered_moves if best_score - scores[move] <= 2.0
    ]
    return random.choice(top_moves), scores

  def _legal_moves(
    self, my_head, my_body, other_snakes, food_set, is_constrictor, width, height
  ):
    occupied = self._occupied_cells(my_body, other_snakes)
    own_tail = (my_body[-1]["x"], my_body[-1]["y"])
    own_tail_stacked = self._is_tail_stacked(my_body)

    safe_moves = {}
    for move, pos in self.get_possible_moves(my_head).items():
      point = (pos["x"], pos["y"])
      if not self._in_bounds(point, width, height):
        continue

      ate_food = point in food_set
      can_step_on_tail = self._can_step_on_own_tail(
        point=point,
        own_tail=own_tail,
        own_tail_is_stacked=own_tail_stacked,
        ate_food=ate_food,
        is_constrictor=is_constrictor,
      )

      if point in occupied and not can_step_on_tail:
        continue

      safe_moves[move] = pos

    return safe_moves

  def _occupied_cells(self, my_body, other_snakes):
    occupied = {(segment["x"], segment["y"]) for segment in my_body}
    for snake in other_snakes:
      occupied |= {(segment["x"], segment["y"]) for segment in snake["body"]}
    return occupied

  def _simulation_blocked(self, future_body, other_snakes, food_set, is_constrictor):
    blocked = {(segment["x"], segment["y"]) for segment in future_body}

    if not is_constrictor and not self._is_tail_stacked(future_body):
      my_tail = future_body[-1]
      blocked.discard((my_tail["x"], my_tail["y"]))

    for snake in other_snakes:
      for segment in snake["body"]:
        blocked.add((segment["x"], segment["y"]))

      if is_constrictor:
        continue

      if self._enemy_can_grow_this_turn(snake, food_set):
        continue

      if self._is_tail_stacked(snake["body"]):
        continue

      enemy_tail = snake["body"][-1]
      blocked.discard((enemy_tail["x"], enemy_tail["y"]))

    return blocked

  def _build_enemy_attack_map(
    self, my_snake, other_snakes, food_set, is_constrictor, width, height
  ):
    occupied = self._occupied_cells(my_snake["body"], other_snakes)
    my_id = my_snake["id"]
    attack_map = {}

    for enemy in other_snakes:
      enemy_len = enemy.get("length", len(enemy["body"]))
      enemy_tail = (enemy["body"][-1]["x"], enemy["body"][-1]["y"])
      enemy_tail_stacked = self._is_tail_stacked(enemy["body"])

      for pos in self.get_possible_moves(enemy["head"]).values():
        point = (pos["x"], pos["y"])
        if not self._in_bounds(point, width, height):
          continue

        can_step_on_enemy_tail = (
          not is_constrictor
          and point == enemy_tail
          and not enemy_tail_stacked
          and not self._enemy_can_grow_this_turn(enemy, food_set)
        )

        if point in occupied and not can_step_on_enemy_tail:
          continue

        # Do not consider impossible overlap directly into my own occupied body except head swap possibilities.
        if point in {
          (segment["x"], segment["y"])
          for segment in my_snake["body"]
          if my_snake["id"] == my_id
        }:
          continue

        previous = attack_map.get(point)
        if previous is None or enemy_len > previous:
          attack_map[point] = enemy_len

    return attack_map

  def _future_body(self, current_body, next_head, ate_food, is_constrictor):
    next_body = [next_head]
    next_body.extend(current_body)

    if is_constrictor or ate_food:
      return next_body

    next_body.pop()
    return next_body

  def _can_step_on_own_tail(
    self, point, own_tail, own_tail_is_stacked, ate_food, is_constrictor
  ):
    if is_constrictor:
      return False
    if ate_food:
      return False
    if own_tail_is_stacked:
      return False
    return point == own_tail

  def _is_tail_stacked(self, body):
    if len(body) < 2:
      return False
    return body[-1]["x"] == body[-2]["x"] and body[-1]["y"] == body[-2]["y"]

  def _enemy_can_grow_this_turn(self, snake, food_set):
    head = snake["head"]
    for dx, dy in self.DIRECTIONS.values():
      if (head["x"] + dx, head["y"] + dy) in food_set:
        return True
    return False

  def _nearest_food_distance(self, start, foods, blocked, width, height):
    if not foods:
      return None

    targets = {(food["x"], food["y"]) for food in foods}
    queue = deque([(start, 0)])
    seen = {start}

    while queue:
      point, distance = queue.popleft()
      if point in targets:
        return distance

      for neighbor in self._neighbors(point):
        if neighbor in seen:
          continue
        if not self._in_bounds(neighbor, width, height):
          continue
        if neighbor in blocked and neighbor not in targets:
          continue
        seen.add(neighbor)
        queue.append((neighbor, distance + 1))

    return None

  def _path_distance(self, start, goal, blocked, width, height):
    queue = deque([(start, 0)])
    seen = {start}

    while queue:
      point, distance = queue.popleft()
      if point == goal:
        return distance

      for neighbor in self._neighbors(point):
        if neighbor in seen:
          continue
        if not self._in_bounds(neighbor, width, height):
          continue
        if neighbor in blocked and neighbor != goal:
          continue
        seen.add(neighbor)
        queue.append((neighbor, distance + 1))

    return None

  def _flood_fill_count(self, start, blocked, width, height):
    queue = deque([start])
    seen = {start}

    while queue:
      point = queue.popleft()
      for neighbor in self._neighbors(point):
        if neighbor in seen:
          continue
        if not self._in_bounds(neighbor, width, height):
          continue
        if neighbor in blocked:
          continue
        seen.add(neighbor)
        queue.append(neighbor)

    return len(seen)

  def _open_neighbor_count(self, start, blocked, width, height):
    count = 0
    for neighbor in self._neighbors(start):
      if not self._in_bounds(neighbor, width, height):
        continue
      if neighbor in blocked:
        continue
      count += 1
    return count

  def _next_turn_option_count(self, future_body, blocked, width, height):
    if not future_body:
      return 0

    next_head = future_body[0]
    count = 0
    for pos in self.get_possible_moves(next_head).values():
      point = (pos["x"], pos["y"])
      if not self._in_bounds(point, width, height):
        continue
      if point in blocked:
        continue
      count += 1
    return count

  def _revisit_penalty(self, point):
    penalty = 0.0
    for index, old_point in enumerate(reversed(self.recent_heads), start=1):
      if old_point != point:
        continue
      penalty += max(0.0, 18.0 - index * 2.0)
    return penalty

  def _territory_control_score(self, my_start, enemy_starts, blocked, width, height):
    if not enemy_starts:
      return 0

    my_distances = self._distance_map(my_start, blocked, width, height)
    enemy_maps = [
      self._distance_map(start, blocked, width, height) for start in enemy_starts
    ]

    score = 0
    for x in range(width):
      for y in range(height):
        point = (x, y)
        if point in blocked:
          continue

        my_distance = my_distances.get(point)
        if my_distance is None:
          continue

        enemy_best = None
        for enemy_map in enemy_maps:
          enemy_distance = enemy_map.get(point)
          if enemy_distance is None:
            continue
          if enemy_best is None or enemy_distance < enemy_best:
            enemy_best = enemy_distance

        if enemy_best is None or my_distance < enemy_best:
          score += 1
        elif enemy_best < my_distance:
          score -= 1

    return score

  def _distance_map(self, start, blocked, width, height):
    queue = deque([(start, 0)])
    distances = {start: 0}

    while queue:
      point, distance = queue.popleft()
      for neighbor in self._neighbors(point):
        if neighbor in distances:
          continue
        if not self._in_bounds(neighbor, width, height):
          continue
        if neighbor in blocked:
          continue
        distances[neighbor] = distance + 1
        queue.append((neighbor, distance + 1))

    return distances

  def _neighbors(self, point):
    for dx, dy in self.DIRECTIONS.values():
      yield (point[0] + dx, point[1] + dy)

  def _manhattan(self, a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

  def _in_bounds(self, point, width, height):
    return 0 <= point[0] < width and 0 <= point[1] < height

  def _fallback_move(self, head, width, height):
    for move, pos in self.get_possible_moves(head).items():
      point = (pos["x"], pos["y"])
      if self._in_bounds(point, width, height):
        return move
    return "up"
