import argparse
import hashlib
import json
from pathlib import Path

class DatasetCurator:
  def __init__(self, input_file: str, output_file: str, min_turn: int = 6, late_turn: int = 20, max_safe_options: int = 2, min_score: int = 3,):
    self.input_file = Path(input_file)
    self.output_file = Path(output_file)
    self.min_turn = min_turn
    self.late_turn = late_turn
    self.max_safe_options = max_safe_options
    self.min_score = min_score

  def _safe_options_count(self, row: dict):
    history = row.get("history", {})
    for item in history.get("data", []):
      if item.get("function") == "get_possible_moves":
        return len(item.get("safe_positions", {}))
    return None

  def _state_hash(self, row: dict):
    board = row.get("game_board", {})
    snakes = board.get("snakes", [])

    snakes_key = []
    for snake in snakes:
      snakes_key.append(
        (
          snake.get("id"),
          snake.get("health"),
          tuple(
            (seg.get("x"), seg.get("y")) for seg in snake.get("body", [])
          ),
        )
      )

    key = {
      "width": board.get("width"),
      "height": board.get("height"),
      "snakes": sorted(snakes_key),
      "food": sorted((f.get("x"), f.get("y")) for f in board.get("food", [])),
      "hazards": sorted(
        (h.get("x"), h.get("y")) for h in board.get("hazards", [])
      ),
    }
    raw = json.dumps(key, sort_keys=True, separators=(",", ":"))
    return hashlib.sha1(raw.encode("utf-8")).hexdigest()

  def _score(self, row: dict):
    score = 0
    turn = int(row.get("turn", 0))
    safe_options = self._safe_options_count(row)
    snakes = row.get("game_board", {}).get("snakes", [])
    opponents = max(0, len(snakes) - 1)

    if turn >= self.late_turn:
      score += 2
    if safe_options is not None and safe_options <= self.max_safe_options:
      score += 3
    if opponents >= 1:
      score += 1

    return score, safe_options

  def curate(self):
    self.output_file.parent.mkdir(parents=True, exist_ok=True)

    total = 0
    kept = 0
    skipped_turn = 0
    skipped_quality = 0
    skipped_duplicate = 0
    seen_states = set()

    with self.input_file.open("r", encoding="utf-8") as src:
      with self.output_file.open("w", encoding="utf-8") as dst:
        for line in src:
          if not line.strip():
            continue

          total += 1
          row = json.loads(line)

          if not row.get("is_good_move", False):
            skipped_quality += 1
            continue

          if int(row.get("turn", 0)) < self.min_turn:
            skipped_turn += 1
            continue

          quality_score, safe_options = self._score(row)
          if quality_score < self.min_score:
            skipped_quality += 1
            continue

          state_key = self._state_hash(row)
          dedupe_key = (state_key, row.get("move"))
          if dedupe_key in seen_states:
            skipped_duplicate += 1
            continue
          seen_states.add(dedupe_key)

          compact_row = {
            "game_id": row.get("game_id"),
            "turn": row.get("turn"),
            "move": row.get("move"),
            "game_type": row.get("game_type"),
            "quality_score": quality_score,
            "safe_options": safe_options,
            "game_board": row.get("game_board"),
          }
          dst.write(json.dumps(compact_row, ensure_ascii=False) + "\n")
          kept += 1

    return {
      "total_rows": total,
      "kept_rows": kept,
      "skipped_turn": skipped_turn,
      "skipped_quality": skipped_quality,
      "skipped_duplicate": skipped_duplicate,
      "output_file": str(self.output_file),
    }

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Create curated best-moves dataset")
  parser.add_argument("--input", required=True, help="Input JSONL file")
  parser.add_argument("--output", required=True, help="Output JSONL file")
  parser.add_argument("--min-turn", type=int, default=6)
  parser.add_argument("--late-turn", type=int, default=20)
  parser.add_argument("--max-safe-options", type=int, default=2)
  parser.add_argument("--min-score", type=int, default=3)
  args = parser.parse_args()

  report = DatasetCurator(
    input_file=args.input,
    output_file=args.output,
    min_turn=args.min_turn,
    late_turn=args.late_turn,
    max_safe_options=args.max_safe_options,
    min_score=args.min_score,
  ).curate()
  print(json.dumps(report, indent=2))