snake-python/server/database/compact/codec.py

"""Byte-level encodings for compact gameplay storage.

All encodings are lossless. Values that do not fit the compact form (for
example coordinates outside 0..255) fall back to a JSON payload, marked by the
leading format byte, so unusual maps can never corrupt stored data.
"""

from datetime import datetime, timezone
import json

MOVES = ("up", "down", "left", "right")
_MOVE_CODES = {move: index + 1 for index, move in enumerate(MOVES)}

_FORMAT_JSON = 0
_FORMAT_COMPACT = 1

# Body step codes relative to the previous segment.
_SAME, _UP, _DOWN, _LEFT, _RIGHT, _ABSOLUTE = range(6)
_STEPS = {(0, 0): _SAME, (0, 1): _UP, (0, -1): _DOWN, (-1, 0): _LEFT, (1, 0): _RIGHT}
_STEP_DELTAS = {code: delta for delta, code in _STEPS.items()}

def encode_move(move:str|None) -> int|None:
  return _MOVE_CODES.get(move) if move is not None else None

def decode_move(code:int|None) -> str|None:
  if not code:
    return None
  return MOVES[int(code) - 1]

def encode_latency(latency:object) -> int|None:
  if latency is None or latency == "":
    return None
  try:
    value = int(str(latency))
  except ValueError:
    return None
  return value if -32768 <= value <= 32767 else None

def decode_latency(value:int|None) -> str|None:
  return None if value is None else str(value)

def _point(point:object) -> tuple[int, int]|None:
  if not isinstance(point, dict):
    return None
  x, y = point.get("x"), point.get("y")
  if type(x) is not int or type(y) is not int:
    return None
  return x, y

def _fits_byte(*values:int) -> bool:
  return all(0 <= value <= 255 for value in values)

def _json_fallback(value:object) -> bytes:
  return bytes([_FORMAT_JSON]) + json.dumps(value, separators=(",", ":")).encode()

def encode_body(body:list) -> bytes:
  """Head as absolute byte pair, then one step code per following segment."""
  points = [_point(segment) for segment in body]
  if any(point is None for point in points) or not all(_fits_byte(*point) for point in points):
    return _json_fallback(body)

  output = bytearray([_FORMAT_COMPACT])
  previous = None
  for x, y in points:
    if previous is None:
      output += bytes((x, y))
    else:
      step = _STEPS.get((x - previous[0], y - previous[1]))
      output += bytes((step,)) if step is not None else bytes((_ABSOLUTE, x, y))
    previous = (x, y)
  return bytes(output)

def decode_body(data:bytes|memoryview|None) -> list[dict]:
  if not data:
    return []
  data = bytes(data)
  if data[0] == _FORMAT_JSON:
    return json.loads(data[1:])

  body = []
  index = 1
  if len(data) >= 3:
    x, y = data[1], data[2]
    body.append({"x": x, "y": y})
    index = 3
  while index < len(data):
    code = data[index]
    if code == _ABSOLUTE:
      x, y = data[index + 1], data[index + 2]
      index += 3
    else:
      dx, dy = _STEP_DELTAS[code]
      x, y = x + dx, y + dy
      index += 1
    body.append({"x": x, "y": y})
  return body

def encode_points(points:list|None) -> bytes|None:
  """Food/hazard lists as byte pairs. ``None`` means "unchanged"."""
  if points is None:
    return None
  pairs = [_point(point) for point in points]
  if any(pair is None for pair in pairs) or not all(_fits_byte(*pair) for pair in pairs):
    return _json_fallback(points)
  output = bytearray([_FORMAT_COMPACT])
  for x, y in pairs:
    output += bytes((x, y))
  return bytes(output)

def decode_points(data:bytes|memoryview|None) -> list[dict]|None:
  if data is None:
    return None
  data = bytes(data)
  if not data:
    return []
  if data[0] == _FORMAT_JSON:
    return json.loads(data[1:])
  return [{"x": data[i], "y": data[i + 1]} for i in range(1, len(data) - 1, 2)]

def flatten_points(points:list[dict]) -> list:
  """Compact JSON form used inside replay blobs: [x0, y0, x1, y1, ...]."""
  pairs = [_point(point) for point in points]
  if any(pair is None for pair in pairs):
    return [points]
  return [value for pair in pairs for value in pair]

def unflatten_points(values:list) -> list[dict]:
  if len(values) == 1 and isinstance(values[0], list):
    return values[0]
  return [{"x": values[i], "y": values[i + 1]} for i in range(0, len(values), 2)]

def to_epoch_us(value:datetime|str|None) -> int|None:
  if value is None:
    return None
  if isinstance(value, str):
    value = datetime.fromisoformat(value.replace("Z", "+00:00"))
  if value.tzinfo is None:
    value = value.replace(tzinfo=timezone.utc)
  delta = value - datetime(1970, 1, 1, tzinfo=timezone.utc)
  return (delta.days * 86_400 + delta.seconds) * 1_000_000 + delta.microseconds

def from_epoch_us(value:int|None) -> datetime|None:
  if value is None:
    return None
  seconds, micros = divmod(int(value), 1_000_000)
  return datetime.fromtimestamp(seconds, tz=timezone.utc).replace(microsecond=micros)