import unittest
from time import perf_counter
from unittest import mock

from snakes.engine.survival_search import CompactSurvivalSearch

from snakes import SnakeBuilder, get_snake_version
from snakes.strategies.prism import PrismBattleSnake_GPT_5_6_Sol
from snakes.strategies.zenith import ZenithBattleSnake_ClaudeOpus5_5

LOSS = -100_000.0
WIN = 100_000.0

class TestZenithRegistration(unittest.TestCase):

  def test_api_name_and_version_are_exposed(self):
    snake = ZenithBattleSnake_ClaudeOpus5_5()

    self.assertEqual(snake.name, "ZenithBattleSnake")
    self.assertEqual(snake.version, "1.2.0")
    self.assertEqual(get_snake_version("ZenithBattleSnake_ClaudeOpus5_5"), "1.2.0")
    self.assertIsInstance(snake, PrismBattleSnake_GPT_5_6_Sol)
    self.assertIsInstance(
      SnakeBuilder.build("ZenithBattleSnake_ClaudeOpus5_5"),
      ZenithBattleSnake_ClaudeOpus5_5,
    )

class TestZenithVerdicts(unittest.TestCase):

  def test_proven_loss_is_vetoed_for_best_surviving_heuristic(self):
    scores = {"up": 300.0, "left": 120.0, "right": 90.0}
    verdicts = {"up": (LOSS, 3), "left": (-40.0, 3), "right": (25.0, 3)}

    self.assertEqual(ZenithBattleSnake_ClaudeOpus5_5._apply_verdicts("up", scores, verdicts), "left")

  def test_proven_win_overrides_heuristic(self):
    scores = {"up": 300.0, "left": 120.0}
    verdicts = {"up": (15.0, 2), "left": (WIN, 2)}

    self.assertEqual(ZenithBattleSnake_ClaudeOpus5_5._apply_verdicts("up", scores, verdicts), "left")

  def test_keeps_choice_when_every_move_loses(self):
    scores = {"up": 300.0, "left": 120.0}
    verdicts = {"up": (LOSS, 2), "left": (LOSS - 1, 2)}

    self.assertEqual(ZenithBattleSnake_ClaudeOpus5_5._apply_verdicts("up", scores, verdicts), "up")

  def test_keeps_choice_without_verdict(self):
    scores = {"up": 300.0, "left": 120.0}
    self.assertEqual(ZenithBattleSnake_ClaudeOpus5_5._apply_verdicts("up", scores, {}), "up")

  def test_verdict_search_resolves_every_candidate(self):
    snake = ZenithBattleSnake_ClaudeOpus5_5()
    # Longer enemy can also reach (2, 5): stepping there loses the head-to-head.
    my_body = [{"x": 1, "y": 5}, {"x": 0, "y": 5}]
    enemy = {"body": [{"x": 3, "y": 5}, {"x": 4, "y": 5}, {"x": 5, "y": 5}], "health": 100}
    safe_moves = {"up": {"x": 1, "y": 6}, "right": {"x": 2, "y": 5}}

    verdicts = snake._duel_verdicts(
      safe_moves=safe_moves, my_body=my_body, enemy=enemy, my_health=100,
      food_set=set(), hazard_set=set(), hazard_damage=15, hazard_count={},
      previous_hazard_set=set(), width=11, height=11,
      deadline=perf_counter() + 1.0,
    )

    self.assertEqual(set(verdicts), {"up", "right"})
    self.assertLessEqual(verdicts["right"][0], LOSS / 2)
    self.assertGreater(verdicts["up"][0], LOSS / 2)
    self.assertEqual(
      snake._apply_verdicts("right", {"right": 50.0, "up": 10.0}, verdicts), "up",
    )

  def test_verdict_search_stops_when_out_of_time(self):
    snake = ZenithBattleSnake_ClaudeOpus5_5()

    verdicts = snake._duel_verdicts(
      safe_moves={"up": {"x": 0, "y": 2}}, my_body=[{"x": 0, "y": 1}],
      enemy={"body": [{"x": 4, "y": 4}]}, my_health=100,
      food_set=set(), hazard_set=set(), hazard_damage=15, hazard_count={},
      previous_hazard_set=set(), width=5, height=5,
      deadline=perf_counter() - 0.001,
    )

    self.assertEqual(verdicts, {})

def _rollout_depth(snake, remaining_ms: float) -> int:
  """Depth requested by a 3-snake rollout with ``remaining_ms`` left."""
  snake._get_bb(11, 11)
  snake._survival_search_context = None
  enemies = [
    {"body": [{"x": 8, "y": 8}, {"x": 8, "y": 9}, {"x": 8, "y": 10}]},
    {"body": [{"x": 2, "y": 8}, {"x": 2, "y": 9}, {"x": 2, "y": 10}]},
  ]
  with mock.patch.object(CompactSurvivalSearch, "search_selected", return_value=0.0) as search:
    snake._future_rollout_bonus(
      move="up", safe_moves={"up": {"x": 5, "y": 3}},
      my_body=[{"x": 5, "y": 2}, {"x": 5, "y": 1}, {"x": 5, "y": 0}],
      other_snakes=enemies, food_set=set(), is_constrictor=False,
      width=11, height=11, enemy_can_grow={},
      deadline=perf_counter() + remaining_ms / 1000,
    )
  return search.call_args.kwargs["depth"]

class TestZenithRolloutDepth(unittest.TestCase):

  def test_gate_is_lowered_by_the_extra_timeout_reserve(self):
    zenith = ZenithBattleSnake_ClaudeOpus5_5()
    prism = PrismBattleSnake_GPT_5_6_Sol()
    extra = zenith._get_timeout_buffer_ms() - prism._get_timeout_buffer_ms()

    self.assertEqual(zenith._rollout_deepen_ms, prism._rollout_deepen_ms - extra)

  def test_three_snake_rollout_deepens_after_same_elapsed_time_as_prism(self):
    # Zenith's budget is 50 ms shorter, so 220 ms left equals 270 ms for Prism.
    self.assertEqual(_rollout_depth(ZenithBattleSnake_ClaudeOpus5_5(), 220), 3)
    self.assertEqual(_rollout_depth(PrismBattleSnake_GPT_5_6_Sol(), 270), 3)
    self.assertEqual(_rollout_depth(PrismBattleSnake_GPT_5_6_Sol(), 220), 2)

  def test_low_budget_still_caps_rollout_depth(self):
    self.assertEqual(_rollout_depth(ZenithBattleSnake_ClaudeOpus5_5(), 150), 2)

class TestZenithVerdictBudget(unittest.TestCase):

  def test_verdict_pass_skips_when_only_reserve_remains(self):
    snake = ZenithBattleSnake_ClaudeOpus5_5()

    verdicts = snake._duel_verdicts(
      safe_moves={"up": {"x": 0, "y": 2}, "right": {"x": 1, "y": 1}},
      my_body=[{"x": 0, "y": 1}], enemy={"body": [{"x": 4, "y": 4}]},
      my_health=100, food_set=set(), hazard_set=set(), hazard_damage=15,
      hazard_count={}, previous_hazard_set=set(), width=5, height=5,
      deadline=perf_counter() + 0.020,
    )

    self.assertEqual(verdicts, {})

if __name__ == "__main__":
  unittest.main()