import assert from "node:assert/strict";
import { readFileSync, readdirSync } from "node:fs";
import { describe, it } from "node:test";
import {
  ARENA_ENVIRONMENTS,
  ARENA_MANIFESTS,
  ARENA_SOURCE_HASHES,
  ArenaScenarioRegistry,
  arenaEnvironmentIds,
  rollout,
  type ArenaManifest,
} from "../../index.ts";

const ARENA_DIR = new URL("../../arena/", import.meta.url);

describe("ARENA_ENVIRONMENTS is the missing half of the catalogue", () => {
  it("has a key for every manifest id, and no key that is not one", () => {
    // Before this existed, `ARENA_MANIFESTS` described environments a harness
    // had no supported way to instantiate: a caller handed "drive-101-v1" off a
    // config file kept its own switch, which is a copy of this catalogue
    // maintained outside the package and wrong the day a sixth env lands.
    const manifestIds = ARENA_MANIFESTS.map((manifest: ArenaManifest) => manifest.id).sort();
    assert.deepEqual(Object.keys(ARENA_ENVIRONMENTS).sort(), manifestIds);
    assert.deepEqual([...arenaEnvironmentIds()].sort(), manifestIds);
    assert.deepEqual(Object.keys(ARENA_SOURCE_HASHES).sort(), manifestIds);
  });

  it("returns an object satisfying the whole ArenaEnvironment shape, keyed to its own id", () => {
    for (const [id, factory] of Object.entries(ARENA_ENVIRONMENTS)) {
      const environment = factory();
      assert.equal(environment.manifest.id, id);
      for (const member of ["reset", "step", "snapshot", "restore", "trace", "replay"] as const) {
        assert.equal(typeof environment[member], "function", `${id}.${member}`);
      }
    }
  });

  it("returns a fresh instance every call, because an episode is state", () => {
    for (const [id, factory] of Object.entries(ARENA_ENVIRONMENTS)) {
      const first = factory();
      const second = factory();
      assert.notEqual(first, second, id);
      first.reset(3, { split: "train" });
      // The second must still be un-reset: two rollouts in flight through the
      // registry must not be stepping each other's episode.
      assert.throws(() => second.step({}), /must be reset/, id);
    }
  });

  it("drives every environment through the shared rollout by id alone", () => {
    for (const id of arenaEnvironmentIds()) {
      const environment = ARENA_ENVIRONMENTS[id]!();
      const result = rollout(environment, () => ({}), { seed: 12, maxSteps: 24 });
      assert.equal(result.steps, 24, id);
      assert.equal(result.final.info.envId, id);
      assert.ok(Number.isFinite(result.total), id);
      // Cut short rather than ended, and both flags say so.
      assert.equal(result.final.terminated || result.final.truncated, false, id);
    }
  });

  it("clamps a rollout budget to the manifest and refuses a budget of nothing", () => {
    const environment = ARENA_ENVIRONMENTS["office-nav-v1"]!();
    const capped = rollout(environment, () => ({}), { seed: 1, maxSteps: 10_000 });
    assert.ok(capped.steps <= environment.manifest.maxSteps);
    assert.throws(
      () => rollout(ARENA_ENVIRONMENTS["office-nav-v1"]!(), () => ({}), { maxSteps: 0 }),
      /at least one step/,
    );
  });
});

describe("scenario selection is bound to the id, not to definition order", () => {
  interface Parameters extends Record<string, number> {
    marker: number;
  }

  function registryOf(ids: readonly string[]): ArenaScenarioRegistry<Parameters> {
    return new ArenaScenarioRegistry<Parameters>(
      "selection-fixture",
      ids.map((id, index) => ({
        id,
        split: id.startsWith("dev") ? ("dev" as const) : ("train" as const),
        parameters: { marker: index },
      })),
      (parameters) => ({ ...parameters }),
    );
  }

  const SEEDS = Array.from({ length: 400 }, (_, index) => index * 7919 + 3);

  it("keeps every seed on the scenario it had when one is inserted in the middle", () => {
    // The trap this replaces: `candidates[seed % candidates.length]` binds a
    // seed to an *array position*, so inserting a scenario silently remaps
    // every seed past it. Nothing fails; the numbers in a results table just
    // quietly stop meaning what they meant.
    const before = registryOf(["train-a", "train-b", "train-c", "dev-a"]);
    const after = registryOf(["train-a", "train-inserted", "train-b", "train-c", "dev-a"]);

    let moved = 0;
    for (const seed of SEEDS) {
      const was = before.resolve(seed, { split: "train" }).id;
      const now = after.resolve(seed, { split: "train" }).id;
      if (now === "train-inserted") continue;
      assert.equal(now, was, `seed ${seed}`);
      moved += 1;
    }
    // And the new scenario genuinely wins some seeds, or the assertion above is
    // vacuous rather than reassuring.
    assert.ok(moved < SEEDS.length, "the inserted scenario must win some seeds");
    assert.ok(moved > SEEDS.length * 0.5, "it must not win most of them either");
  });

  it("is unaffected by reordering the literal at all", () => {
    const declared = registryOf(["train-a", "train-b", "train-c", "dev-a"]);
    const shuffled = registryOf(["train-c", "train-a", "dev-a", "train-b"]);
    for (const seed of SEEDS) {
      assert.equal(
        shuffled.resolve(seed, { split: "train" }).id,
        declared.resolve(seed, { split: "train" }).id,
        `seed ${seed}`,
      );
    }
  });

  it("spreads seeds across the split rather than parking them on one scenario", () => {
    const registry = registryOf(["train-a", "train-b", "train-c", "dev-a"]);
    const counts = new Map<string, number>();
    for (const seed of SEEDS) {
      const id = registry.resolve(seed, { split: "train" }).id;
      counts.set(id, (counts.get(id) ?? 0) + 1);
    }
    assert.equal(counts.size, 3);
    for (const [id, count] of counts) assert.ok(count > SEEDS.length / 6, `${id}=${count}`);
  });

  it("bumped every shipped manifest's version, because selection changed under them", () => {
    // The other half of the fix. A selection change that nothing recorded is
    // the same silent remap; `version` is what a snapshot, a trace and a
    // results table are pinned to, so it moves when selection does.
    for (const manifest of ARENA_MANIFESTS) {
      const expected = manifest.id === "studio-ops-v1" ? 1 : 2;
      assert.equal(manifest.version, expected, manifest.id);
    }
  });
});

describe("the arena boundary holds", () => {
  const SOURCES = readdirSync(ARENA_DIR)
    .filter((name) => name.endsWith(".ts"))
    .map((name) => ({ name, text: readFileSync(new URL(name, ARENA_DIR), "utf8") }));

  it("has sources to check", () => {
    assert.ok(SOURCES.length >= 12);
  });

  it("imports no three.js, no scene adapter, no DOM and no network", () => {
    // The executable form of ARENA.md's first paragraph and of the build spec's
    // grep. An arena that reached the renderer would be an arena that cannot be
    // run headless, and a `studio-ops-v1` that imported `engine/flights.ts` for
    // its aircraft would have done exactly that on one line.
    const forbidden =
      /from ["'](three|three\/[^"']*|\.\.\/engine\/(scene|stage|scenekit|flights|atmosphere|world)\.ts|\.\.\/actors\/sceneActor\.ts|\.\.\/interiors\/officeScene\.ts)["']/;
    for (const source of SOURCES) {
      assert.equal(forbidden.test(source.text), false, `${source.name} imports the renderer`);
      assert.equal(/\bdocument\.|\bwindow\.|\bfetch\(/.test(source.text), false, source.name);
    }
  });

  it("declares a simulator pin for every file studio-ops actually wraps", () => {
    // A pin list shorter than the import list is a snapshot surviving a change
    // it should not have survived, so the two are compared rather than trusted.
    const studioOps = SOURCES.find((source) => source.name === "studioOps.ts")!;
    const script = readFileSync(new URL("../../../scripts/check-arena-source-hashes.mjs", import.meta.url), "utf8");
    const pinned = script.slice(script.indexOf('"studio-ops-v1"'));
    for (const match of studioOps.text.matchAll(/from "\.\.\/([^"]+)"/g)) {
      const relative = `src/${match[1]}`;
      assert.ok(pinned.includes(`"${relative}"`), `${relative} is imported but not pinned`);
    }
  });
});
