import asyncio
import unittest
from unittest.mock import AsyncMock

from server.database.benchmark_states import is_postgresql_source, replay_turn_to_state, sample_states

GAME = {
  "game_id": "game-id", "source": "league", "map": "standard",
  "ruleset_name": "standard", "ruleset_version": "v1",
}

def _turn(number:int) -> dict:
  you = {"id": "you-id", "head": {"x": 1, "y": 2}, "body": [{"x": 1, "y": 2}]}
  return {
    "turn": number,
    "board": {"width": 11, "height": 11, "food": [{"x": 5, "y": 5}], "hazards": [], "snakes": [you]},
    "you": you,
  }

class TestBenchmarkStates(unittest.TestCase):
  def test_detects_postgresql_dsn(self):
    self.assertTrue(is_postgresql_source("postgresql://user:pass@example.com/db"))
    self.assertTrue(is_postgresql_source("postgres://user:pass@example.com/db"))
    self.assertFalse(is_postgresql_source("/tmp/gameplay.sqlite3"))

  def test_builds_state_from_replay_turn(self):
    board, metadata = replay_turn_to_state(GAME, _turn(7))
    self.assertEqual(board["food"], [{"x": 5, "y": 5}])
    self.assertEqual(metadata["you"]["id"], "you-id")
    self.assertEqual(metadata["turn"], 7)
    self.assertEqual(metadata["ruleset"]["version"], "v1")

  def test_turn_without_snakes_is_skipped(self):
    self.assertIsNone(replay_turn_to_state(GAME, {"turn": 0, "board": {"snakes": []}, "you": {}}))

  def test_samples_every_stride_turn_from_newest_games(self):
    backend = AsyncMock()
    backend.list_game_ids.return_value = ["old", "new"]
    backend.get_game_replay.side_effect = lambda game_id: {
      "game": {**GAME, "game_id": game_id},
      "turns": [_turn(number) for number in range(4)],
    }

    states = asyncio.run(sample_states(backend, samples=3, stride=3))

    self.assertEqual([(m["game_id"], m["turn"]) for _, m in states], [("new", 3), ("new", 0), ("old", 1)])

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