/** Renderer-independent, JSON-safe reinforcement-learning contract. */

export const ARENA_API_VERSION = "tera.arena/v1" as const;

export type ArenaSplit = "train" | "dev";

export interface ArenaScenarioRequest {
  /** Public splits only. Held-out/private evaluation belongs outside this package. */
  split: ArenaSplit;
  /** Omit to choose deterministically from the requested split using the seed. */
  id?: string;
}

export interface ArenaScenario<P extends object = Record<string, number | string | boolean>> {
  id: string;
  split: ArenaSplit;
  seed: number;
  parameters: P;
  hash: string;
}

export interface ArenaSourceHashes {
  /** Hash of the environment implementation source, pinned and checked in CI. */
  environment: string;
  /** Hash of the renderer-independent controller/plan sources the environment wraps. */
  simulator: string;
}

export interface ArenaInfo {
  apiVersion: typeof ARENA_API_VERSION;
  envId: string;
  envVersion: number;
  envHash: string;
  scenarioId: string;
  scenarioSplit: ArenaSplit;
  scenarioHash: string;
  simulatorHash: string;
  environmentSourceHash: string;
  seed: number;
  step: number;
  maxSteps: number;
  terminalReason: string | null;
  stateChecksum: string;
}

export interface ArenaResetResult<O> {
  observation: O;
  info: ArenaInfo;
}

export interface ArenaStepResult<O, R extends Record<string, number>> {
  observation: O;
  reward: number;
  rewardComponents: R;
  terminated: boolean;
  truncated: boolean;
  info: ArenaInfo;
}

/**
 * How one observation or action field is encoded as numbers.
 *
 * `observationFields` was always a list of names, which tells a reader what the
 * environment observes and tells a trainer nothing it can allocate against: how
 * wide is the vector, which slots are angles, what does `2` mean in the slot
 * holding a phase. Every consumer had to read the environment's source and
 * write the encoder by hand, which is a fork of the observation contract kept
 * in a file this repository cannot see.
 *
 * Four kinds, because four is what the six environments actually contain:
 *
 * - `float` — a quantity, with `low`/`high` as *declared bounds* rather than
 *   guarantees. `flattenObservation` clamps into them; a policy normalises with
 *   them. One slot.
 * - `bool` — one slot, 0 or 1.
 * - `enum` — a closed vocabulary, one-hot over `values`. As many slots as
 *   there are values. A value outside the vocabulary encodes as all zeros,
 *   which reads as "none of these" rather than silently colliding with the
 *   first member.
 * - `id` — an open-vocabulary identifier: a station id, a robot id, a road
 *   name. One slot holding a stable hash in `[0, 1)`. Nothing should try to
 *   *learn* from that slot — it has no order and no metric — but dropping the
 *   field would make the vector and `observationFields` incommensurable, and a
 *   hash at least changes exactly when the identity does, which is enough to
 *   detect "the station I am walking to is not the one I was walking to".
 */
export type ArenaFieldKind = "float" | "bool" | "enum" | "id";

export interface ArenaFieldSpec {
  /** Matches the key in the observation or action object, exactly. */
  name: string;
  kind: ArenaFieldKind;
  /** `float` only. Both required for a float, finite, and `low < high`. */
  low?: number;
  high?: number;
  /** `enum` only. Non-empty and free of duplicates. */
  values?: readonly string[];
  /** Documentation for a reader of the manifest. Never parsed. */
  unit?: string;
}

export interface ArenaManifest {
  apiVersion: typeof ARENA_API_VERSION;
  id: string;
  version: number;
  title: string;
  description: string;
  simulator: string;
  fixedStepSeconds: number;
  maxSteps: number;
  actionFields: readonly string[];
  observationFields: readonly string[];
  /**
   * The same fields as `actionFields`/`observationFields`, in the same order,
   * with an encoding for each.
   *
   * Two lists rather than one because the name lists are the human-readable
   * contract that has been published since v1 and the spaces are the machine
   * one; `src/test/arena/spaces.test.ts` asserts they agree name for name and
   * in order, so the redundancy cannot rot into a disagreement.
   */
  actionSpace: readonly ArenaFieldSpec[];
  observationSpace: readonly ArenaFieldSpec[];
  rewardComponents: Readonly<Record<string, string>>;
  safetyTerminals: readonly string[];
  scenarioIds: Readonly<Record<ArenaSplit, readonly string[]>>;
  baselines: {
    inaction: string;
    scripted: string;
  };
}

export interface ArenaSnapshot<S> {
  apiVersion: typeof ARENA_API_VERSION;
  envId: string;
  envVersion: number;
  envHash: string;
  /**
   * The source pins the snapshot was taken under.
   *
   * `envHash` is a hash of the *manifest*, which is semantics: the id, the
   * version, the field names, the reward vocabulary. It does not move when the
   * physics under it does. Edit the walker's collision epsilon or the device
   * simulator's release rate, leave the manifest alone, and a snapshot taken
   * before the edit restored cleanly afterwards and resumed into a different
   * simulation — silently, mid-episode, with a checksum that still verified
   * because it was recomputed over the new state.
   *
   * `replay()` has always guarded this by comparing the trace envelope's
   * `sourceHashes`; `restore()` did not, and a checkpoint is exactly as
   * dangerous as a trace. So the pins travel in the snapshot too, inside the
   * checksummed core, and `restore()` refuses a snapshot whose simulator or
   * environment source has moved. The cost is that a legitimate source edit
   * invalidates outstanding checkpoints, which is the correct thing for it to
   * do: they were checkpoints of a program that no longer exists.
   */
  sourceHashes: ArenaSourceHashes;
  seed: number;
  scenarioId: string;
  scenarioSplit: ArenaSplit;
  scenarioHash: string;
  step: number;
  cumulativeReward: number;
  terminated: boolean;
  truncated: boolean;
  terminalReason: string | null;
  simulation: S;
  checksum: string;
}

export interface ArenaTraceStep<A, R extends Record<string, number>> {
  index: number;
  action: A;
  reward: number;
  rewardComponents: R;
  terminated: boolean;
  truncated: boolean;
  terminalReason: string | null;
  stateChecksum: string;
}

export interface ArenaTraceEnvelope<A, R extends Record<string, number>> {
  apiVersion: typeof ARENA_API_VERSION;
  envId: string;
  envVersion: number;
  envHash: string;
  scenarioId: string;
  scenarioSplit: ArenaSplit;
  scenarioHash: string;
  sourceHashes: ArenaSourceHashes;
  seed: number;
  initialStateChecksum: string;
  steps: readonly ArenaTraceStep<A, R>[];
  finalStateChecksum: string;
  cumulativeReward: number;
  checksum: string;
}

export interface ArenaReplayResult<O> {
  observation: O;
  steps: number;
  cumulativeReward: number;
  finalStateChecksum: string;
}

export interface ArenaEnvironment<
  A,
  O,
  R extends Record<string, number>,
  S,
> {
  readonly manifest: ArenaManifest;
  reset(seed: number, scenario: string | ArenaScenarioRequest): ArenaResetResult<O>;
  step(action: A): ArenaStepResult<O, R>;
  snapshot(): ArenaSnapshot<S>;
  restore(snapshot: ArenaSnapshot<S>): ArenaResetResult<O>;
  trace(): ArenaTraceEnvelope<A, R>;
  replay(trace: ArenaTraceEnvelope<A, R>): ArenaReplayResult<O>;
}
