import assert from "node:assert/strict";
import { describe, it } from "node:test";
import {
  ARENA_API_VERSION,
  ARENA_MANIFESTS,
  ARENA_SOURCE_HASHES,
  CALIFORNIA_FLIGHT_INACTION,
  CALIFORNIA_FLIGHT_SCENARIOS,
  CROW_NAV_INACTION,
  CROW_NAV_SCENARIOS,
  DRIVE_101_SCENARIOS,
  DRIVE_INACTION,
  OFFICE_NAV_INACTION,
  OFFICE_NAV_SCENARIOS,
  OFFICE_JOBS_INACTION,
  OFFICE_JOBS_SCENARIOS,
  STUDIO_OPS_INACTION,
  STUDIO_OPS_SCENARIOS,
  CaliforniaFlightEnvironment,
  CrowNavEnvironment,
  Drive101Environment,
  OfficeNavEnvironment,
  OfficeJobsEnvironment,
  StudioOpsEnvironment,
  arenaChecksum,
  californiaFlightScriptedBaseline,
  crowNavScriptedBaseline,
  driveScriptedBaseline,
  officeNavScriptedBaseline,
  officeJobsScriptedBaseline,
  rollout,
  studioOpsScriptedBaseline,
  type ArenaEnvironment,
  type ArenaManifest,
  type ArenaScenarioRegistry,
  type ArenaStepResult,
} from "../index.ts";

type NumericRewards = Record<string, number>;
type AnyEnvironment = ArenaEnvironment<unknown, unknown, NumericRewards, unknown>;

interface EnvironmentCase {
  name: string;
  create(): AnyEnvironment;
  registry: ArenaScenarioRegistry<object>;
  inaction: unknown;
  scripted(observation: unknown): unknown;
  successReason: string;
  /**
   * Where the documented inaction action ends up.
   *
   * `"max-steps"` for the five environments whose only clock is the step cap.
   * `studio-ops-v1` has a second one — a departure the studio's car has to be
   * ready for — and standing still misses it, which is an *outcome* and so sets
   * `terminated`. The field exists so the contract test can keep asserting the
   * thing that actually matters (`truncated` is set by the step cap and by
   * nothing else) rather than being weakened to accommodate the sixth case.
   */
  inactionReason: string;
}

const CASES: EnvironmentCase[] = [
  {
    name: "drive",
    create: () => new Drive101Environment() as AnyEnvironment,
    registry: DRIVE_101_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: DRIVE_INACTION,
    scripted: () => driveScriptedBaseline(),
    successReason: "goal",
    inactionReason: "max-steps",
  },
  {
    name: "office",
    create: () => new OfficeNavEnvironment() as AnyEnvironment,
    registry: OFFICE_NAV_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: OFFICE_NAV_INACTION,
    scripted: (observation) => officeNavScriptedBaseline(
      observation as Parameters<typeof officeNavScriptedBaseline>[0],
    ),
    successReason: "goal",
    inactionReason: "max-steps",
  },
  {
    name: "office-jobs",
    create: () => new OfficeJobsEnvironment() as AnyEnvironment,
    registry: OFFICE_JOBS_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: OFFICE_JOBS_INACTION,
    scripted: (observation) => officeJobsScriptedBaseline(
      observation as Parameters<typeof officeJobsScriptedBaseline>[0],
    ),
    successReason: "job-complete",
    inactionReason: "max-steps",
  },
  {
    name: "crow",
    create: () => new CrowNavEnvironment() as AnyEnvironment,
    registry: CROW_NAV_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: CROW_NAV_INACTION,
    scripted: (observation) => crowNavScriptedBaseline(
      observation as Parameters<typeof crowNavScriptedBaseline>[0],
    ),
    successReason: "goal",
    inactionReason: "max-steps",
  },
  {
    name: "flight",
    create: () => new CaliforniaFlightEnvironment() as AnyEnvironment,
    registry: CALIFORNIA_FLIGHT_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: CALIFORNIA_FLIGHT_INACTION,
    scripted: (observation) => californiaFlightScriptedBaseline(
      observation as Parameters<typeof californiaFlightScriptedBaseline>[0],
    ),
    successReason: "goal",
    inactionReason: "max-steps",
  },
  {
    name: "studio-ops",
    create: () => new StudioOpsEnvironment() as AnyEnvironment,
    registry: STUDIO_OPS_SCENARIOS as ArenaScenarioRegistry<object>,
    inaction: STUDIO_OPS_INACTION,
    scripted: (observation) => studioOpsScriptedBaseline(
      observation as Parameters<typeof studioOpsScriptedBaseline>[0],
    ),
    successReason: "job-complete",
    inactionReason: "departure-missed",
  },
];

/**
 * One episode, through the package's own `rollout`.
 *
 * This function used to be a hand-written loop, and so did the two baseline
 * proofs below and every example in ARENA.md — four copies of the same nine
 * lines, agreeing by luck. `rollout` is now the one copy, and running the
 * contract suite through it means a regression in the shared loop fails here
 * rather than in a consumer's trainer.
 */
function run(
  entry: EnvironmentCase,
  scenarioId: string,
  seed: number,
  policy: (observation: unknown) => unknown,
): { total: number; final: ArenaStepResult<unknown, NumericRewards> } {
  const result = rollout(entry.create(), (observation) => policy(observation), {
    seed,
    scenario: scenarioId,
  });
  return { total: result.total, final: result.final };
}

describe("arena contract and manifests", () => {
  it("exports six versioned renderer-independent manifests with disjoint public splits", () => {
    assert.deepEqual(ARENA_MANIFESTS.map((manifest: ArenaManifest) => manifest.id), [
      "drive-101-v1", "office-nav-v1", "office-jobs-v1", "crow-nav-v1", "california-flight-v1",
      "studio-ops-v1",
    ]);
    for (const manifest of ARENA_MANIFESTS) {
      assert.equal(manifest.apiVersion, ARENA_API_VERSION);
      assert.ok(manifest.version >= 1 && manifest.maxSteps > 0);
      assert.ok(manifest.safetyTerminals.length > 0);
      assert.ok(Object.keys(manifest.rewardComponents).length >= 6);
      assert.equal(
        manifest.scenarioIds.train.some((id) => manifest.scenarioIds.dev.includes(id)),
        false,
      );
      assert.match(ARENA_SOURCE_HASHES[manifest.id]!.environment, /^sha256:[0-9a-f]{64}$/);
      assert.match(ARENA_SOURCE_HASHES[manifest.id]!.simulator, /^sha256:[0-9a-f]{64}$/);
    }
  });

  it("selects and materializes scenarios deterministically by split, id and seed", () => {
    for (const entry of CASES) {
      const a = entry.create().reset(0x1234abcd, { split: "train" });
      const b = entry.create().reset(0x1234abcd, { split: "train" });
      assert.deepEqual(a, b, entry.name);
      assert.equal(a.info.scenarioSplit, "train");
      assert.match(a.info.scenarioHash, /^fnv1a64:[0-9a-f]{16}$/);
      const different = entry.create().reset(0x1234abce, a.info.scenarioId);
      assert.notEqual(different.info.scenarioHash, a.info.scenarioHash);
      assert.throws(
        () => entry.create().reset(1, { split: "dev", id: entry.registry.ids("train")[0] }),
        /unknown/,
      );
    }
  });

  it("sanitizes non-finite and out-of-range actions before state or reward hashing", () => {
    const malformed = {
      throttle: Number.POSITIVE_INFINITY,
      brake: Number.NaN,
      steering: -99,
      handbrake: false,
    };
    const first = new Drive101Environment();
    const second = new Drive101Environment();
    first.reset(4, "train-us101-ventura");
    second.reset(4, "train-us101-ventura");
    const a = first.step(malformed);
    const b = second.step({ throttle: 0, brake: 0, steering: -1, handbrake: false });
    assert.deepEqual(a, b);
    assert.ok(Number.isFinite(a.reward));
  });

  it("uses terminated for outcomes, truncated only for max-step exhaustion", () => {
    for (const entry of CASES) {
      const id = entry.registry.ids("train")[0]!;
      const idle = run(entry, id, 5, () => entry.inaction).final;
      // `truncated` is set by the step cap and by nothing else, in every
      // environment. An inaction outcome that is not `max-steps` is a genuine
      // terminal and must therefore set `terminated` instead — which is the
      // whole distinction this test exists to pin.
      assert.equal(idle.info.terminalReason, entry.inactionReason, entry.name);
      assert.equal(idle.truncated, entry.inactionReason === "max-steps", entry.name);
      assert.equal(idle.terminated, entry.inactionReason !== "max-steps", entry.name);

      const scripted = run(entry, id, 5, entry.scripted).final;
      assert.equal(scripted.terminated, true, entry.name);
      assert.equal(scripted.truncated, false, entry.name);
      assert.equal(scripted.info.terminalReason, entry.successReason, entry.name);
    }
  });
});

describe("arena baseline proofs", () => {
  it("keeps inaction below zero and proves a positive scripted completion on every public scenario", () => {
    for (const entry of CASES) {
      for (const definition of entry.registry.definitions) {
        for (const seed of [1, 0xdecafbad]) {
          const idle = run(entry, definition.id, seed, () => entry.inaction);
          assert.ok(idle.total < 0, `${entry.name}/${definition.id}/${seed} inaction=${idle.total}`);
          assert.notEqual(idle.final.info.terminalReason, entry.successReason);

          const scripted = run(entry, definition.id, seed, entry.scripted);
          assert.equal(
            scripted.final.info.terminalReason,
            entry.successReason,
            `${entry.name}/${definition.id}/${seed}`,
          );
          assert.ok(scripted.total > 0, `${entry.name}/${definition.id}/${seed}=${scripted.total}`);
          assert.ok(scripted.total > idle.total);
        }
      }
    }
  });

  it("exposes reachable safety terminals rather than reward-only safety labels", () => {
    const drive = new Drive101Environment();
    drive.reset(2, "train-us101-ventura");
    let driveReason: string | null = null;
    for (let index = 0; index < drive.manifest.maxSteps; index += 1) {
      const result = drive.step({ throttle: 1, brake: 0, steering: 1, handbrake: false });
      driveReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(driveReason, "guardrail-contact");

    const office = new OfficeNavEnvironment();
    office.reset(2, "train-hangar-crossing");
    let officeReason: string | null = null;
    for (let index = 0; index < office.manifest.maxSteps; index += 1) {
      const result = office.step({ x: 0, z: 1 });
      officeReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(officeReason, "collision-stall");

    const officeJobs = new OfficeJobsEnvironment();
    officeJobs.reset(2, "train-la-directory-inspection");
    let officeJobsReason: string | null = null;
    for (let index = 0; index < 8; index += 1) {
      const result = officeJobs.step({ x: 0, z: 0, interact: true });
      officeJobsReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(officeJobsReason, "wrong-interaction-limit");

    const collisionJobs = new OfficeJobsEnvironment();
    collisionJobs.reset(2, "dev-sf-parcel-delivery");
    let collisionJobsReason: string | null = null;
    for (let index = 0; index < 80; index += 1) {
      const result = collisionJobs.step({ x: 0, z: 1, interact: false });
      collisionJobsReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(collisionJobsReason, "collision-stall");

    const crow = new CrowNavEnvironment();
    crow.reset(2, "train-east-crosswind");
    let crowReason: string | null = null;
    for (let index = 0; index < crow.manifest.maxSteps; index += 1) {
      const result = crow.step({ forward: 1, turn: 0, pitch: 0, climb: 1, glide: false });
      crowReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(crowReason, "flight-envelope-contact");

    const flight = new CaliforniaFlightEnvironment();
    flight.reset(2, "train-la-east-leg");
    let flightReason: string | null = null;
    for (let index = 0; index < flight.manifest.maxSteps; index += 1) {
      const result = flight.step({ throttle: 1, yaw: 0, pitch: 1, roll: 0 });
      flightReason = result.info.terminalReason;
      if (result.terminated) break;
    }
    assert.equal(flightReason, "flight-envelope-contact");
  });
});

describe("arena snapshot, trace and replay", () => {
  it("restores and replays every SF/LA office-job scenario", () => {
    for (const definition of OFFICE_JOBS_SCENARIOS.definitions) {
      const environment = new OfficeJobsEnvironment();
      let observation = environment.reset(311, definition.id).observation;
      for (let index = 0; index < 10; index += 1) {
        observation = environment.step(officeJobsScriptedBaseline(observation)).observation;
      }
      const checkpoint = environment.snapshot();
      const action = officeJobsScriptedBaseline(observation);
      const expected = environment.step(action);
      environment.restore(checkpoint);
      assert.deepEqual(environment.step(action), expected, definition.id);

      const traced = new OfficeJobsEnvironment();
      observation = traced.reset(311, definition.id).observation;
      for (let index = 0; index < 50; index += 1) {
        const result = traced.step(officeJobsScriptedBaseline(observation));
        observation = result.observation;
        if (result.terminated) break;
      }
      const trace = traced.trace();
      const replay = new OfficeJobsEnvironment().replay(trace);
      assert.equal(replay.finalStateChecksum, trace.finalStateChecksum, definition.id);
      assert.equal(replay.cumulativeReward, trace.cumulativeReward, definition.id);
    }
  });

  it("restores each simulator bit-for-bit and preserves the next transition", () => {
    for (const entry of CASES) {
      const environment = entry.create();
      let observation = environment.reset(91, entry.registry.ids("dev")[0]!).observation;
      for (let index = 0; index < 12; index += 1) {
        observation = environment.step(entry.scripted(observation)).observation;
      }
      const checkpoint = environment.snapshot();
      const action = entry.scripted(observation);
      const expected = environment.step(action);
      const restored = environment.restore(checkpoint);
      assert.equal(restored.info.step, checkpoint.step);
      const actual = environment.step(action);
      assert.deepEqual(actual, expected, entry.name);
    }
  });

  it("replays checksummed action traces to the exact final state", () => {
    for (const entry of CASES) {
      const environment = entry.create();
      let observation = environment.reset(144, entry.registry.ids("train")[0]!).observation;
      for (let index = 0; index < 30; index += 1) {
        const result = environment.step(entry.scripted(observation));
        observation = result.observation;
        if (result.terminated) break;
      }
      const trace = environment.trace();
      assert.equal(arenaChecksum({ ...trace, checksum: undefined }), trace.checksum);
      const replay = entry.create().replay(trace);
      assert.equal(replay.finalStateChecksum, trace.finalStateChecksum, entry.name);
      assert.equal(replay.cumulativeReward, trace.cumulativeReward, entry.name);
      assert.equal(replay.steps, trace.steps.length, entry.name);
    }
  });

  it("rejects tampered snapshots and traces before applying them", () => {
    const environment = new Drive101Environment();
    environment.reset(7, "train-us101-ventura");
    environment.step(driveScriptedBaseline());
    const snapshot = environment.snapshot();
    assert.throws(
      () => environment.restore({ ...snapshot, cumulativeReward: snapshot.cumulativeReward + 1 }),
      /checksum mismatch/,
    );
    const trace = environment.trace();
    assert.throws(
      () => new Drive101Environment().replay({ ...trace, cumulativeReward: trace.cumulativeReward + 1 }),
      /checksum mismatch/,
    );

    const invalidCore = { ...snapshot, step: -1 };
    const { checksum: _oldChecksum, ...invalidWithoutChecksum } = invalidCore;
    assert.throws(
      () => environment.restore({
        ...invalidWithoutChecksum,
        checksum: arenaChecksum(invalidWithoutChecksum),
      }),
      /episode state is invalid/,
    );

    const invalidFrames = trace.steps.map((frame, index) => ({
      ...frame,
      index: index === 0 ? 2 : frame.index,
    }));
    const { checksum: _traceChecksum, ...invalidTraceCore } = { ...trace, steps: invalidFrames };
    assert.throws(
      () => new Drive101Environment().replay({
        ...invalidTraceCore,
        checksum: arenaChecksum(invalidTraceCore),
      }),
      /frame 1 is invalid/,
    );

    const jobs = new OfficeJobsEnvironment();
    jobs.reset(17, "train-sf-display-inspection");
    const jobsSnapshot = jobs.snapshot();
    const invalidSimulation = structuredClone(jobsSnapshot.simulation);
    invalidSimulation.activity.robots[0]!.battery = 9;
    const { checksum: _jobsChecksum, ...jobsCore } = { ...jobsSnapshot, simulation: invalidSimulation };
    assert.throws(
      () => jobs.restore({ ...jobsCore, checksum: arenaChecksum(jobsCore) }),
      /invalid/,
    );
  });

  it("refuses post-terminal stepping until reset", () => {
    const environment = new OfficeNavEnvironment();
    let observation = environment.reset(5, "train-galley-aisle").observation;
    for (;;) {
      const result = environment.step(officeNavScriptedBaseline(observation));
      observation = result.observation;
      if (result.terminated) break;
    }
    assert.throws(() => environment.step(OFFICE_NAV_INACTION), /episode is complete/);
  });
});
