"""In-memory records shared by both backends, the replay blob and migrations.

A ``TurnRecord`` is already in compact form: ``thinking`` has derivable fields
stripped, and ``food``/``hazards`` are ``None`` when unchanged since the
previous stored turn.
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone

from .codec import encode_latency
from .thinking import derived_values, slim_thinking

@dataclass(frozen=True)
class SnakeIdentity:
  snake_id:str
  name:str|None
  is_you:bool
  customizations:dict = field(default_factory=dict)

@dataclass(frozen=True)
class SnakeState:
  snake_id:str
  turn:int
  health:int|None
  latency:int|None
  body:list

@dataclass(frozen=True)
class TurnRecord:
  turn:int
  observed_at:datetime|None
  move:str|None
  thinking:object
  food:list|None
  hazards:list|None

@dataclass(frozen=True)
class PointsState:
  """Last stored food/hazards, used to store only changes."""
  turn:int
  food:list
  hazards:list

@dataclass
class GameBundle:
  game:dict
  snakes:list[SnakeIdentity]
  turns:list[TurnRecord]
  states:list[SnakeState]

def you_view(states:list[SnakeState], you_id:str|None) -> dict|None:
  state = next((s for s in states if s.snake_id == you_id), None)
  if state is None:
    return None
  return {
    "head": state.body[0] if state.body else None,
    "health": state.health,
    "length": len(state.body),
  }

def compact_turn(
  turn:int,
  observed_at:datetime|None,
  move:str|None,
  thinking:object,
  food:list,
  hazards:list,
  states:list[SnakeState],
  you_id:str|None,
  width:int|None,
  height:int|None,
  previous:PointsState|None,
) -> tuple[TurnRecord, PointsState]:
  derived = derived_values(turn, move, you_view(states, you_id), food, width, height, len(states))
  unchanged = previous is not None and previous.turn == turn - 1
  record = TurnRecord(
    turn=turn,
    observed_at=observed_at,
    move=move,
    thinking=slim_thinking(thinking, derived),
    food=None if unchanged and previous.food == food else food,
    hazards=None if unchanged and previous.hazards == hazards else hazards,
  )
  return record, PointsState(turn, food, hazards)

def records_from_game_state(
  game_state:dict,
  my_move:str|None,
  thinking:object,
  previous:PointsState|None,
) -> tuple[TurnRecord, list[SnakeState], list[SnakeIdentity], PointsState]:
  board = game_state.get("board", {})
  you_id = game_state.get("you", {}).get("id")
  turn = int(game_state.get("turn", 0))
  identities = []
  states = []
  for snake in board.get("snakes", []):
    snake_id = snake.get("id")
    if snake_id is None:
      continue
    identities.append(SnakeIdentity(
      snake_id=snake_id,
      name=snake.get("name"),
      is_you=snake_id == you_id,
      customizations=snake.get("customizations") or {},
    ))
    states.append(SnakeState(
      snake_id=snake_id,
      turn=turn,
      health=snake.get("health"),
      latency=encode_latency(snake.get("latency")),
      body=list(snake.get("body", [])),
    ))
  record, points = compact_turn(
    turn, datetime.now(timezone.utc), my_move, thinking,
    list(board.get("food", [])), list(board.get("hazards", [])),
    states, you_id, board.get("width"), board.get("height"), previous,
  )
  return record, states, identities, points

def merge_records(
  base_turns:list[TurnRecord],
  base_states:list[SnakeState],
  turns:list[TurnRecord],
  states:list[SnakeState],
) -> tuple[list[TurnRecord], list[SnakeState]]:
  """Combine a packed replay with rows written after it was packed."""
  merged_turns = {record.turn: record for record in base_turns}
  merged_turns.update({record.turn: record for record in turns})
  merged_states = {(state.snake_id, state.turn): state for state in base_states}
  merged_states.update({(state.snake_id, state.turn): state for state in states})
  return (
    sorted(merged_turns.values(), key=lambda record: record.turn),
    sorted(merged_states.values(), key=lambda state: (state.turn, state.snake_id)),
  )