from pathlib import Path
from typing import Any
import os

from server.dataset.DatasetIO import DatasetIO

class RLBootstrapDataset:
  def __init__(self):
    self.enabled = self._env_bool("RL_BOOTSTRAP_ENABLED", default=False)
    self.min_base_rows = self._env_int("RL_MIN_BASE_ROWS", default=5000)
    self.base_dataset_path = Path(os.getenv("RL_BASE_DATASET", "data/dataset/best_moves.jsonl"))
    self.output_path = Path(os.getenv("RL_BOOTSTRAP_OUTPUT", "data/dataset/rl_bootstrap.jsonl"))
    self.max_bytes = int(float(os.getenv("RL_BOOTSTRAP_MAX_MB", "50")) * 1024 * 1024)
    self.needs_more_data = False

  @staticmethod
  def _env_bool(name:str, default:bool=False) -> bool:
    value = os.getenv(name)
    if value is None:
      return default
    return value.lower() in {"1", "true", "yes", "on"}

  @staticmethod
  def _env_int(name:str, default:int) -> int:
    value = os.getenv(name)
    if value is None:
      return default
    try:
      return int(value)
    except ValueError:
      return default

  def refresh_state(self):
    if not self.enabled:
      self.needs_more_data = False
      return

    base_rows = DatasetIO.count_jsonl_rows(self.base_dataset_path)
    self.needs_more_data = base_rows < self.min_base_rows

  def record_sample(self, game_data:Any, move:str, safe_moves:dict[str, dict[str, int]], reason:str, scores:dict[str, float]|None=None):
    if not self.enabled or not self.needs_more_data:
      return

    try:
      self.output_path.parent.mkdir(parents=True, exist_ok=True)
      row = {
        "source": "best_battlesnake_bootstrap",
        "game_id": getattr(game_data, "id", None),
        "turn": game_data.get_turn(),
        "move": move,
        "safe_moves": list(safe_moves.keys()),
        "reason": reason,
        "game_board": game_data.get_game_board_as_dict(),
      }
      if scores:
        row["scores"] = {k: round(v, 5) for k, v in scores.items()}

      DatasetIO.append_jsonl_row(self.output_path, row)
      DatasetIO.rotate_and_gzip_if_size_reached(self.output_path, self.max_bytes)
    except Exception:
      return