import unittest

import tempfile, json, gzip
from pathlib import Path

from server.dataset.DatasetExporter import DatasetExporter

class TestDatasetExporter(unittest.TestCase):
  def test_export_jsonl(self):
    with tempfile.TemporaryDirectory() as tmp:
      input_dir = Path(tmp) / "data"
      output_file = Path(tmp) / "out" / "dataset.jsonl"
      game_file = input_dir / "game-1.json"
      game_file.parent.mkdir(parents=True, exist_ok=True)

      game_payload = {
        "dataset": {
          "game": {"id": "g-1", "map": "standard", "type": {"name": "duel"}},
          "snake": {"type": "BestBattleSnake"},
          "samples": [
            {
              "turn": 1,
              "move": "up",
              "is_good_move": True,
              "game_board": {"width": 11, "height": 11},
              "history": {"data": []},
            }
          ],
        }
      }
      game_file.write_text(json.dumps(game_payload), encoding="utf-8")

      report = DatasetExporter(str(input_dir), str(output_file)).export_jsonl()

      self.assertEqual(report["games_scanned"], 1)
      self.assertEqual(report["samples_exported"], 1)
      self.assertTrue(output_file.exists())

      lines = output_file.read_text(encoding="utf-8").strip().splitlines()
      self.assertEqual(len(lines), 1)
      first = json.loads(lines[0])
      self.assertEqual(first["game_id"], "g-1")
      self.assertEqual(first["move"], "up")
      self.assertTrue(first["is_good_move"])

  def test_export_jsonl_gz(self):
    with tempfile.TemporaryDirectory() as tmp:
      input_dir = Path(tmp) / "data"
      output_file = Path(tmp) / "out" / "dataset.jsonl.gz"
      game_file = input_dir / "game-1.json"
      game_file.parent.mkdir(parents=True, exist_ok=True)

      game_payload = {
        "dataset": {
          "game": {"id": "g-1", "map": "standard", "type": {"name": "duel"}},
          "snake": {"type": "BestBattleSnake"},
          "samples": [
            {
              "turn": 1,
              "move": "up",
              "is_good_move": True,
              "game_board": {"width": 11, "height": 11},
              "history": {"data": []},
            }
          ],
        }
      }
      game_file.write_text(json.dumps(game_payload), encoding="utf-8")

      report = DatasetExporter(str(input_dir), str(output_file)).export_jsonl()

      self.assertEqual(report["games_scanned"], 1)
      self.assertEqual(report["samples_exported"], 1)
      self.assertTrue(output_file.exists())

      with gzip.open(output_file, "rt", encoding="utf-8") as handle:
        lines = handle.read().strip().splitlines()
      self.assertEqual(len(lines), 1)
      first = json.loads(lines[0])
      self.assertEqual(first["game_id"], "g-1")

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