"""Deterministic scenario corpus for the local snake arena."""

from __future__ import annotations

from copy import deepcopy

from tests.bench_best_battle_snake import build_game_state

def _state(payload: dict, scenario: str, index: int) -> tuple[dict, dict]:
  payload["game"]["id"] = f"arena-{scenario}-{index}"
  return payload["board"], {
    "game_id": payload["game"]["id"],
    "source": "custom",
    "map": payload["game"].get("map", "standard"),
    "ruleset": payload["game"]["ruleset"],
    "turn": payload["turn"],
    "you": payload["you"],
    "scenario": scenario,
  }

def _standard_duel(index: int) -> dict:
  payload = build_game_state()
  payload["turn"] = 20 + index
  payload["board"]["food"] = [
    {"x": 1 + index % 3, "y": 9},
    {"x": 9, "y": 1 + (index // 3) % 3},
  ]
  return payload

def _hazard_duel(index: int) -> dict:
  payload = _standard_duel(index)
  payload["you"]["health"] = 38 + index % 12
  payload["board"]["snakes"][0]["health"] = payload["you"]["health"]
  hazard_x = 5 + index % 2
  payload["board"]["hazards"] = [
    {"x": hazard_x, "y": y} for y in range(1, 10) if y != 5
  ]
  return payload

def _multiplayer(index: int) -> dict:
  payload = _standard_duel(index)
  third = {
    "id": "enemy-2",
    "name": "enemy-2",
    "health": 65,
    "length": 5,
    "head": {"x": 2, "y": 8},
    "body": [
      {"x": 2, "y": 8}, {"x": 2, "y": 9}, {"x": 2, "y": 10},
      {"x": 1, "y": 10}, {"x": 0, "y": 10},
    ],
  }
  payload["board"]["snakes"].append(third)
  return payload

def _constrictor(index: int) -> dict:
  payload = _multiplayer(index)
  payload["game"]["ruleset"] = deepcopy(payload["game"]["ruleset"])
  payload["game"]["ruleset"]["name"] = "constrictor"
  payload["board"]["food"] = []
  return payload

def _cramped_duel(index: int) -> dict:
  payload = _standard_duel(index)
  payload["board"]["width"] = 7
  payload["board"]["height"] = 7
  mine = {
    "id": "me", "name": "me", "health": 72, "length": 7,
    "head": {"x": 2, "y": 3},
    "body": [
      {"x": 2, "y": 3}, {"x": 2, "y": 2}, {"x": 2, "y": 1},
      {"x": 1, "y": 1}, {"x": 1, "y": 2}, {"x": 1, "y": 3},
      {"x": 1, "y": 4},
    ],
  }
  enemy = {
    "id": "enemy", "name": "enemy", "health": 72, "length": 7,
    "head": {"x": 4, "y": 3},
    "body": [
      {"x": 4, "y": 3}, {"x": 4, "y": 2}, {"x": 4, "y": 1},
      {"x": 5, "y": 1}, {"x": 5, "y": 2}, {"x": 5, "y": 3},
      {"x": 5, "y": 4},
    ],
  }
  payload["you"] = mine
  payload["board"]["snakes"] = [mine, enemy]
  payload["board"]["food"] = [{"x": 3, "y": 5 + index % 2}]
  payload["board"]["hazards"] = []
  return payload

SCENARIOS = {
  "duel": _standard_duel,
  "hazard": _hazard_duel,
  "multi": _multiplayer,
  "constrictor": _constrictor,
  "cramped": _cramped_duel,
}

def synthetic_states(count: int, scenarios: list[str] | None = None) -> list[tuple[dict, dict]]:
  selected = scenarios or list(SCENARIOS)
  unknown = set(selected) - set(SCENARIOS)
  if unknown:
    raise ValueError(f"Unknown arena scenarios: {', '.join(sorted(unknown))}")
  states: list[tuple[dict, dict]] = []
  for index in range(count):
    scenario = selected[index % len(selected)]
    states.append(_state(SCENARIOS[scenario](index), scenario, index))
  return states