from snakes.TemplateSnake import TemplateSnake

from queue import PriorityQueue, Queue
import random

class AStarSnake(TemplateSnake):
  def avoid_my_body(self, my_body, possible_moves: dict) -> list:
    """
    my_body: Set of tuples representing x/y coordinates for every segment of a Battlesnake.
            e.g. {(0, 0), (1, 0), (2, 0)}
    possible_moves: Dictionary of moves to pick from, with coordinates as tuples.
            e.g. {"up": (0, 1), "down": (0, -1), "left": (-1, 0), "right": (1, 0)}

    return: The dictionary of remaining possible_moves, with the moves leading to self-collision removed
    """
    remove = []
    for direction, location in possible_moves.items():
      if location in my_body:
        remove.append(direction)

    for direction in remove:
      del possible_moves[direction]

    return possible_moves

  def avoid_walls(self, board_width: int, board_height: int, possible_moves: dict):
    remove = []
    for direction, location in possible_moves.items():
      x_out_range = (location[0] < 0 or location[0] == board_width)
      y_out_range = (location[1] < 0 or location[1] == board_height)
      if x_out_range or y_out_range:
        remove.append(direction)

    for direction in remove:
      del possible_moves[direction]

    return possible_moves

  def avoid_snakes(self, snakes: list, possible_moves: dict):
    remove = []
    for snake in snakes:
      for direction, location in possible_moves.items():
        if location in snake["body"]:
          remove.append(direction)

    remove = set(remove)
    for direction in remove:
      del possible_moves[direction]

    return possible_moves

  def get_target_close(self, foods: list, my_head: tuple):
    if len(foods) == 0:
      return None
    return min(foods, key=lambda food: abs(food["x"] - my_head[0]) + abs(food["y"] - my_head[1]))

  def flood_fill(self, board_width: int, board_height: int, my_body: set):
    """
    Perform Flood Fill to identify safe areas on the board.
    """
    visited = set()
    safe_cells = set()

    # Define directions (up, down, left, right)
    directions = [(0, 1), (0, -1), (-1, 0), (1, 0)]

    # Perform Flood Fill from each cell not occupied by the snake's body
    for x in range(board_width):
      for y in range(board_height):
        if (x, y) not in my_body and (x, y) not in visited:
          # Start Flood Fill from this cell
          q = Queue()
          q.put((x, y))
          visited.add((x, y))
          safe_cells.add((x, y))

          # Continue Flood Fill until the queue is empty
          while not q.empty():
            current_cell = q.get()
            for dx, dy in directions:
              new_cell = (current_cell[0] + dx, current_cell[1] + dy)
              if 0 <= new_cell[0] < board_width and 0 <= new_cell[1] < board_height:
                if new_cell not in my_body and new_cell not in visited:
                  q.put(new_cell)
                  visited.add(new_cell)
                  safe_cells.add(new_cell)

    return safe_cells

  def choose_move(self, data: dict) -> str:
    my_head = (data["you"]["head"]["x"], data["you"]["head"]["y"])
    my_body = {(part["x"], part["y"]) for part in data["you"]["body"]}
    board_height = data["board"]["height"]
    board_width = data["board"]["width"]
    foods = data["board"]["food"]

    # Perform Flood Fill to identify safe areas on the board
    safe_cells = self.flood_fill(board_width, board_height, my_body)

    # Find the nearest food located in a safe area using A* algorithm
    def heuristic(a, b):
      return abs(a[0] - b[0]) + abs(a[1] - b[1])

    def a_star(start, goal:set=()):
      open_set = Queue()
      open_set.put(start)
      came_from = {}
      g_score = {start: 0}

      while not open_set.empty():
        current = open_set.get()

        if current == goal:
          path = []
          while current in came_from:
            path.append(current)
            current = came_from[current]
          return path[::-1][0]

        for dx, dy in [(0, 1), (0, -1), (-1, 0), (1, 0)]:
          new_cell = (current[0] + dx, current[1] + dy)
          if new_cell in safe_cells:
            tentative_g_score = g_score[current] + 1
            if new_cell not in g_score or tentative_g_score < g_score[new_cell]:
              came_from[new_cell] = current
              g_score[new_cell] = tentative_g_score
              open_set.put(new_cell)

      return None

    try:
      nearest_food = min(foods, key=lambda food: heuristic(my_head, (food["x"], food["y"])))
      target_position = (nearest_food["x"], nearest_food["y"])
      move_target = a_star(my_head, target_position)
    except ValueError:
      # TODO: What to do when no food is available?
      # - Avoid own body and other snakes
      # - Flut fill?
      move_target = a_star(my_head, safe_cells)
      print(move_target)

    # Choose the next move based on the path obtained from A* algorithm
    if move_target:
      dx = move_target[0] - my_head[0]
      dy = move_target[1] - my_head[1]

      self.add_to_history({"my_head": my_head, "my_body": tuple(my_body), "target": move_target, "target_position": target_position, "nearest_food": nearest_food, "dx": dx, "dy": dy})
      if dx == 1:
        return "right"
      elif dx == -1:
        return "left"
      elif dy == 1:
        return "up"
      elif dy == -1:
        return "down"

    # If no safe path to food is found, choose a random move
    random_move = random.choice(["up", "down", "left", "right"])

    try:
      self.add_to_history({"my_head": my_head, "my_body": tuple(my_body), "target": move_target, "target_position": target_position, "nearest_food": nearest_food, "random_move": random_move})
    except UnboundLocalError:
      self.add_to_history({"my_head": my_head, "my_body": tuple(my_body), "target": move_target, "random_move": random_move})

    return random_move