/**
 * One episode, run to its end.
 *
 * Nine lines of loop that every consumer of this package had written for
 * itself: the test file, both baseline proofs, and every example in ARENA.md.
 * Four copies of a loop is four places for the step budget to be off by one,
 * for `truncated` to be checked and `terminated` not to be, or for the returned
 * total to be the environment's cumulative reward in one copy and the caller's
 * own running sum in another. They agree today; nothing made them.
 *
 * It is deliberately thin. A trainer does not want a framework here — it wants
 * the loop it was going to write, exported once so that a regression in it is a
 * regression in one place. Anything richer (vectorised environments, batching,
 * an action buffer) belongs in the trainer, which knows things this package
 * cannot: how many workers there are, what device the policy is on, and whether
 * an episode is worth finishing.
 */

import type {
  ArenaEnvironment,
  ArenaScenarioRequest,
  ArenaStepResult,
} from "./types.ts";

/**
 * What a policy is, from this package's point of view.
 *
 * The step index is passed as well as the observation because a *scripted*
 * baseline sometimes needs it — the studio-ops targeted policies switch
 * behaviour at a known step to reach a particular terminal — and because a
 * policy that wants only the observation can ignore a second argument for free.
 * It is the step that is about to be taken, counting from zero.
 */
export type ArenaPolicy<A, O> = (observation: O, step: number) => A;

export interface RolloutOptions {
  /** Defaults to 0, which is a legal seed and a boring one. */
  seed?: number;
  /** An exact public scenario id, or a split request. Defaults to `train`. */
  scenario?: string | ArenaScenarioRequest;
  /**
   * Stop after this many steps even if the episode has not ended.
   *
   * Clamped to the manifest's own cap, because a caller cannot extend an
   * episode past the point where the environment sets `truncated` — asking for
   * more steps than the manifest allows is a mistake worth silently correcting
   * rather than a request worth honouring. Below the cap it is a genuine early
   * cut, and `final.terminated`/`final.truncated` will both be false, which is
   * how a caller tells "I stopped it" from "it ended".
   */
  maxSteps?: number;
}

export interface RolloutResult<O, R extends Record<string, number>> {
  /** The sum of every step's total reward. */
  total: number;
  /** How many steps were actually taken. */
  steps: number;
  /** The last transition, carrying the terminal reason and final observation. */
  final: ArenaStepResult<O, R>;
}

export function rollout<A, O, R extends Record<string, number>, S>(
  environment: ArenaEnvironment<A, O, R, S>,
  policy: ArenaPolicy<A, O>,
  options: RolloutOptions = {},
): RolloutResult<O, R> {
  const budget = Math.min(
    environment.manifest.maxSteps,
    options.maxSteps ?? environment.manifest.maxSteps,
  );
  if (!Number.isSafeInteger(budget) || budget < 1) {
    throw new RangeError("arena rollout needs a budget of at least one step");
  }

  let observation = environment.reset(
    options.seed ?? 0,
    options.scenario ?? { split: "train" },
  ).observation;

  let total = 0;
  let steps = 0;
  let final: ArenaStepResult<O, R> | undefined;
  for (let step = 0; step < budget; step += 1) {
    const result = environment.step(policy(observation, step));
    observation = result.observation;
    total += result.reward;
    steps += 1;
    final = result;
    if (result.terminated || result.truncated) break;
  }
  // Unreachable while `budget >= 1`, and asserted rather than assumed because
  // the alternative is a non-null assertion that would go stale the moment the
  // loop above grows a `continue`.
  if (!final) throw new Error("arena rollout took no steps");
  return { total, steps, final };
}
