import { arenaChecksum, quantizeForChecksum } from "./checksum.ts";
import type { ArenaScenarioRegistry } from "./scenarios.ts";
import {
  ARENA_API_VERSION,
  type ArenaEnvironment,
  type ArenaInfo,
  type ArenaManifest,
  type ArenaReplayResult,
  type ArenaResetResult,
  type ArenaScenario,
  type ArenaScenarioRequest,
  type ArenaSnapshot,
  type ArenaSourceHashes,
  type ArenaStepResult,
  type ArenaTraceEnvelope,
  type ArenaTraceStep,
} from "./types.ts";

export interface SimulationTransition<O, R extends Record<string, number>> {
  observation: O;
  rewardComponents: R;
  terminated?: boolean;
  terminalReason?: string;
}

/** Shared episode, checksum, snapshot and replay semantics for concrete environments. */
export abstract class BaseArenaEnvironment<
  A,
  O,
  R extends Record<string, number>,
  S,
  P extends object,
> implements ArenaEnvironment<A, O, R, S> {
  abstract readonly manifest: ArenaManifest;
  protected abstract readonly registry: ArenaScenarioRegistry<P>;
  protected abstract readonly sourceHashes: ArenaSourceHashes;

  private scenario: ArenaScenario<P> | null = null;
  private stepIndex = 0;
  private cumulativeReward = 0;
  private terminated = false;
  private truncated = false;
  private terminalReason: string | null = null;
  private initialStateChecksum = "";
  private frames: ArenaTraceStep<A, R>[] = [];

  protected abstract resetSimulation(scenario: ArenaScenario<P>): O;
  protected abstract normalizeAction(action: A): A;
  protected abstract advanceSimulation(action: A): SimulationTransition<O, R>;
  protected abstract simulationSnapshot(): S;
  protected abstract restoreSimulation(snapshot: S): O;

  reset(seed: number, request: string | ArenaScenarioRequest): ArenaResetResult<O> {
    this.scenario = this.registry.resolve(seed, request);
    this.stepIndex = 0;
    this.cumulativeReward = 0;
    this.terminated = false;
    this.truncated = false;
    this.terminalReason = null;
    this.frames = [];
    const observation = this.resetSimulation(this.scenario);
    this.initialStateChecksum = arenaChecksum(this.statePayload());
    return { observation, info: this.info() };
  }

  step(rawAction: A): ArenaStepResult<O, R> {
    this.requireReset();
    if (this.terminated || this.truncated) {
      throw new Error("arena episode is complete; call reset before step");
    }
    const action = this.normalizeAction(rawAction);
    const transition = this.advanceSimulation(action);
    const values = Object.values(transition.rewardComponents);
    if (values.some((value) => !Number.isFinite(value))) {
      throw new Error("arena reward components must be finite");
    }
    const reward = values.reduce((sum, value) => sum + value, 0);
    this.stepIndex += 1;
    this.cumulativeReward += reward;
    this.terminated = transition.terminated === true;
    this.truncated = !this.terminated && this.stepIndex >= this.manifest.maxSteps;
    this.terminalReason = transition.terminalReason ?? (this.truncated ? "max-steps" : null);
    const stateChecksum = arenaChecksum(this.statePayload());
    this.frames.push({
      index: this.stepIndex,
      action: structuredClone(action),
      reward,
      rewardComponents: structuredClone(transition.rewardComponents),
      terminated: this.terminated,
      truncated: this.truncated,
      terminalReason: this.terminalReason,
      stateChecksum,
    });
    return {
      observation: transition.observation,
      reward,
      rewardComponents: structuredClone(transition.rewardComponents),
      terminated: this.terminated,
      truncated: this.truncated,
      info: this.info(stateChecksum),
    };
  }

  snapshot(): ArenaSnapshot<S> {
    const scenario = this.requireReset();
    const core = {
      apiVersion: ARENA_API_VERSION,
      envId: this.manifest.id,
      envVersion: this.manifest.version,
      envHash: this.envHash(),
      // Inside the checksummed core, so a snapshot cannot be re-pinned to a
      // different simulator without invalidating itself. See `ArenaSnapshot`.
      sourceHashes: this.sourceHashes,
      seed: scenario.seed,
      scenarioId: scenario.id,
      scenarioSplit: scenario.split,
      scenarioHash: scenario.hash,
      step: this.stepIndex,
      cumulativeReward: this.cumulativeReward,
      terminated: this.terminated,
      truncated: this.truncated,
      terminalReason: this.terminalReason,
      simulation: structuredClone(this.simulationSnapshot()),
    };
    return { ...core, checksum: arenaChecksum(core) };
  }

  restore(snapshot: ArenaSnapshot<S>): ArenaResetResult<O> {
    const { checksum, ...core } = snapshot;
    if (arenaChecksum(core) !== checksum) throw new Error("arena snapshot checksum mismatch");
    if (
      snapshot.apiVersion !== ARENA_API_VERSION || snapshot.envId !== this.manifest.id ||
      snapshot.envVersion !== this.manifest.version || snapshot.envHash !== this.envHash() ||
      arenaChecksum(snapshot.sourceHashes) !== arenaChecksum(this.sourceHashes)
    ) throw new Error("arena snapshot is incompatible with this environment");
    if (
      !Number.isSafeInteger(snapshot.step) || snapshot.step < 0 ||
      snapshot.step > this.manifest.maxSteps || !Number.isFinite(snapshot.cumulativeReward) ||
      typeof snapshot.terminated !== "boolean" || typeof snapshot.truncated !== "boolean" ||
      (snapshot.terminated && snapshot.truncated) ||
      (snapshot.terminalReason !== null && typeof snapshot.terminalReason !== "string")
    ) throw new Error("arena snapshot episode state is invalid");
    const resolved = this.registry.resolve(snapshot.seed, {
      split: snapshot.scenarioSplit,
      id: snapshot.scenarioId,
    });
    if (resolved.hash !== snapshot.scenarioHash) throw new Error("arena scenario hash mismatch");
    this.scenario = resolved;
    this.stepIndex = snapshot.step;
    this.cumulativeReward = snapshot.cumulativeReward;
    this.terminated = snapshot.terminated;
    this.truncated = snapshot.truncated;
    this.terminalReason = snapshot.terminalReason;
    this.frames = [];
    let observation: O;
    try {
      observation = this.restoreSimulation(structuredClone(snapshot.simulation));
    } catch (error) {
      // A concrete environment can reject the simulation payload after the
      // episode bookkeeping above has already been written — and several do,
      // because validating a controller's snapshot means handing it to the
      // controller. Leaving the object in that state was the worst of the three
      // options: the step index, the cumulative reward and the terminal flags
      // would say one thing and the simulators another, and the *next* call
      // would succeed and quietly produce a mixture of two episodes.
      //
      // Rolling back is not available either: `restoreSimulation` has usually
      // rebuilt the simulators from the scenario before it validates, so the
      // state it threw from is not the state it started in. So the environment
      // is marked un-reset, which makes every subsequent call fail loudly until
      // somebody calls `reset` or `restore` again — the same disposition
      // `devices/sim.ts` takes when it refuses a snapshot outright.
      this.scenario = null;
      throw error;
    }
    this.initialStateChecksum = arenaChecksum(this.statePayload());
    return { observation, info: this.info() };
  }

  trace(): ArenaTraceEnvelope<A, R> {
    const scenario = this.requireReset();
    const core = {
      apiVersion: ARENA_API_VERSION,
      envId: this.manifest.id,
      envVersion: this.manifest.version,
      envHash: this.envHash(),
      scenarioId: scenario.id,
      scenarioSplit: scenario.split,
      scenarioHash: scenario.hash,
      sourceHashes: this.sourceHashes,
      seed: scenario.seed,
      initialStateChecksum: this.initialStateChecksum,
      steps: structuredClone(this.frames),
      finalStateChecksum: arenaChecksum(this.statePayload()),
      cumulativeReward: this.cumulativeReward,
    };
    return { ...core, checksum: arenaChecksum(core) };
  }

  replay(trace: ArenaTraceEnvelope<A, R>): ArenaReplayResult<O> {
    const { checksum, ...core } = trace;
    if (arenaChecksum(core) !== checksum) throw new Error("arena trace checksum mismatch");
    if (
      trace.apiVersion !== ARENA_API_VERSION || trace.envId !== this.manifest.id ||
      trace.envVersion !== this.manifest.version ||
      trace.envHash !== this.envHash() || arenaChecksum(trace.sourceHashes) !== arenaChecksum(this.sourceHashes)
    ) throw new Error("arena trace is incompatible with this environment");
    if (
      !Array.isArray(trace.steps) || trace.steps.length > this.manifest.maxSteps ||
      !Number.isFinite(trace.cumulativeReward)
    ) throw new Error("arena trace episode state is invalid");
    let reset = this.reset(trace.seed, { split: trace.scenarioSplit, id: trace.scenarioId });
    if (trace.scenarioHash !== reset.info.scenarioHash || trace.initialStateChecksum !== reset.info.stateChecksum) {
      throw new Error("arena trace initial state mismatch");
    }
    let observation = reset.observation;
    for (let offset = 0; offset < trace.steps.length; offset += 1) {
      const expected = trace.steps[offset]!;
      if (
        expected.index !== offset + 1 || !Number.isFinite(expected.reward) ||
        typeof expected.terminated !== "boolean" || typeof expected.truncated !== "boolean" ||
        (expected.terminated && expected.truncated)
      ) throw new Error(`arena trace frame ${offset + 1} is invalid`);
      const actual = this.step(expected.action);
      observation = actual.observation;
      if (
        actual.info.stateChecksum !== expected.stateChecksum ||
        // Quantised rather than `!==`, for the reason `checksum.ts` sets out at
        // length: a reward is a float, a verifier runs on hardware the producer
        // never saw, and `Math.pow` is not required to be correctly rounded. An
        // exact comparison here would have re-introduced on one line precisely
        // the cross-runtime failure the checksum was hardened against.
        quantizeForChecksum(actual.reward) !== quantizeForChecksum(expected.reward) ||
        actual.terminated !== expected.terminated || actual.truncated !== expected.truncated ||
        arenaChecksum(actual.rewardComponents) !== arenaChecksum(expected.rewardComponents)
      ) throw new Error(`arena trace diverged at step ${expected.index}`);
    }
    const replayed = this.trace();
    if (
      replayed.finalStateChecksum !== trace.finalStateChecksum ||
      quantizeForChecksum(replayed.cumulativeReward) !== quantizeForChecksum(trace.cumulativeReward)
    ) throw new Error("arena trace final state mismatch");
    return {
      observation,
      steps: this.stepIndex,
      cumulativeReward: this.cumulativeReward,
      finalStateChecksum: replayed.finalStateChecksum,
    };
  }

  protected currentScenario(): ArenaScenario<P> {
    return this.requireReset();
  }

  protected currentStep(): number {
    return this.stepIndex;
  }

  private requireReset(): ArenaScenario<P> {
    if (!this.scenario) throw new Error("arena environment must be reset before use");
    return this.scenario;
  }

  private envHash(): string {
    return arenaChecksum(this.manifest);
  }

  private statePayload(): object {
    const scenario = this.requireReset();
    return {
      envId: this.manifest.id,
      envVersion: this.manifest.version,
      seed: scenario.seed,
      scenarioHash: scenario.hash,
      step: this.stepIndex,
      cumulativeReward: this.cumulativeReward,
      terminated: this.terminated,
      truncated: this.truncated,
      terminalReason: this.terminalReason,
      simulation: this.simulationSnapshot(),
    };
  }

  private info(checksum = arenaChecksum(this.statePayload())): ArenaInfo {
    const scenario = this.requireReset();
    return {
      apiVersion: ARENA_API_VERSION,
      envId: this.manifest.id,
      envVersion: this.manifest.version,
      envHash: this.envHash(),
      scenarioId: scenario.id,
      scenarioSplit: scenario.split,
      scenarioHash: scenario.hash,
      simulatorHash: this.sourceHashes.simulator,
      environmentSourceHash: this.sourceHashes.environment,
      seed: scenario.seed,
      step: this.stepIndex,
      maxSteps: this.manifest.maxSteps,
      terminalReason: this.terminalReason,
      stateChecksum: checksum,
    };
  }
}
