"""Pack a finished game's turns into one zstd-compressed blob.

Consecutive turns are nearly identical, so compressing a whole game at once is
far smaller than storing one row per turn and snake.
"""

from compression import zstd
import json

from .codec import (
  decode_move, encode_move, flatten_points, from_epoch_us, to_epoch_us,
  unflatten_points,
)
from .records import SnakeState, TurnRecord

REPLAY_FORMAT = 1
_LEVEL = 10

def _optional_points(points:list|None) -> list|None:
  return None if points is None else flatten_points(points)

def _optional_unflatten(values:list|None) -> list|None:
  return None if values is None else unflatten_points(values)

def pack_replay(turns:list[TurnRecord], states:list[SnakeState]) -> bytes:
  by_snake:dict[str, list] = {}
  for state in sorted(states, key=lambda s: (s.snake_id, s.turn)):
    by_snake.setdefault(state.snake_id, []).append(
      [state.turn, state.health, state.latency, flatten_points(state.body)]
    )
  payload = {
    "t": [
      [
        record.turn, to_epoch_us(record.observed_at), encode_move(record.move),
        record.thinking, _optional_points(record.food), _optional_points(record.hazards),
      ]
      for record in sorted(turns, key=lambda item: item.turn)
    ],
    "s": by_snake,
  }
  raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode()
  return zstd.compress(raw, level=_LEVEL)

def unpack_replay(data:bytes|memoryview, format_version:int=REPLAY_FORMAT) -> tuple[list[TurnRecord], list[SnakeState]]:
  if format_version != REPLAY_FORMAT:
    raise ValueError(f"Unsupported replay format: {format_version}")
  payload = json.loads(zstd.decompress(bytes(data)))
  turns = [
    TurnRecord(
      turn=turn, observed_at=from_epoch_us(observed), move=decode_move(move),
      thinking=thinking, food=_optional_unflatten(food), hazards=_optional_unflatten(hazards),
    )
    for turn, observed, move, thinking, food, hazards in payload["t"]
  ]
  states = [
    SnakeState(
      snake_id=snake_id, turn=turn, health=health, latency=latency,
      body=unflatten_points(body),
    )
    for snake_id, rows in payload["s"].items()
    for turn, health, latency, body in rows
  ]
  return turns, states