from pathlib import Path
from typing import Any
import random, json, os

from server.TrainBattleSnakeAI import MOVES, extract_feature_values
from snakes.TemplateSnake import TemplateSnake

class TrainedBattleSnake(TemplateSnake):
  VERSION = "0.1.0"

  def __init__(self):
    super().__init__()
    self.name = "TrainedBattleSnake"
    self.version = self.VERSION
    self._model_path:Path|None=None
    self._model_data:dict[str, Any]|None=None

  def choose_move(self, game_data) -> str:
    self.game_board = game_data
    self.calculations = []

    safe_positions = self.find_safe_positions(add_to_calculations=True)
    if not safe_positions:
      self.add_to_history({"turn": game_data.get_turn(), "reason": "no_safe_moves"})
      return "up"

    model = self._load_model()
    if not model:
      move = random.choice(list(safe_positions.keys()))
      self.add_to_history({
        "turn": game_data.get_turn(),
        "move": move,
        "reason": "model_missing",
        "safe_moves": list(safe_positions.keys()),
      })
      return move

    row = {
      "turn": game_data.get_turn(),
      "game_board": game_data.get_game_board_as_dict(),
    }
    scores = self._predict_scores(model, row)

    best_safe_move = max(safe_positions.keys(), key=lambda move: scores.get(move, float("-inf")))
    self.add_to_history({
      "turn": game_data.get_turn(),
      "move": best_safe_move,
      "safe_moves": list(safe_positions.keys()),
      "scores": {move: round(scores.get(move, 0.0), 5) for move in MOVES},
    })
    return best_safe_move

  def _load_model(self) -> dict[str, Any] | None:
    env_path = os.getenv("TRAINED_SNAKE_MODEL", "models/battlesnake_softmax_v2.json")
    path = Path(env_path)

    if self._model_path == path and self._model_data is not None:
      return self._model_data

    if not path.exists() or not path.is_file():
      self._model_path = path
      self._model_data = None
      return None

    payload = json.loads(path.read_text(encoding="utf-8"))
    model = payload.get("model")
    if not isinstance(model, dict):
      self._model_path = path
      self._model_data = None
      return None

    self._model_path = path
    self._model_data = model
    return model

  def _predict_scores(self, model:dict[str, Any], row:dict[str, Any]) -> dict[str, float]:
    return self._predict_scores_softmax_v2(model, row)

  def _predict_scores_softmax_v2(self, model:dict[str, Any], row:dict[str, Any]) -> dict[str, float]:
    features = extract_feature_values(row)
    weights = model.get("weights", {})
    bias = model.get("bias", {})
    scores:dict[str, float] = {}

    for move in MOVES:
      move_weights = weights.get(move, {})
      score = float(bias.get(move, 0.0))
      for name, value in features.items():
        score += float(move_weights.get(name, 0.0)) * float(value)
      scores[move] = score
    return scores