import assert from "node:assert/strict";
import { describe, it } from "node:test";
import {
  ARENA_ENVIRONMENTS,
  ARENA_MANIFESTS,
  STUDIO_OPS_INACTION,
  STUDIO_OPS_SCENARIOS,
  StudioOpsEnvironment,
  actionWidth,
  arenaEnvironmentIds,
  arenaFieldWidth,
  arenaManifest,
  flattenAction,
  flattenObservation,
  observationWidth,
  structureAction,
  studioOpsScriptedBaseline,
  type ArenaFieldSpec,
  type StudioOpsAction,
  type ArenaManifest,
  type StudioOpsObservation,
} from "../../index.ts";

describe("every manifest declares a space that matches its field list", () => {
  it("names the same fields, in the same order, on both lists", () => {
    // The two lists are redundant on purpose — names are the contract that has
    // been published since v1, spaces are the machine-readable one — and this
    // is what stops the redundancy rotting into a disagreement.
    for (const manifest of ARENA_MANIFESTS as readonly ArenaManifest[]) {
      assert.deepEqual(
        manifest.observationSpace.map((spec: ArenaFieldSpec) => spec.name),
        [...manifest.observationFields],
        manifest.id,
      );
      assert.deepEqual(
        manifest.actionSpace.map((spec: ArenaFieldSpec) => spec.name),
        [...manifest.actionFields],
        manifest.id,
      );
    }
  });

  it("declares a usable encoding for every field", () => {
    for (const manifest of ARENA_MANIFESTS as readonly ArenaManifest[]) {
      for (const spec of [...manifest.observationSpace, ...manifest.actionSpace]) {
        // `arenaFieldWidth` is where a malformed spec is caught, so calling it
        // over the whole catalogue is the validation pass.
        assert.ok(arenaFieldWidth(spec) >= 1, `${manifest.id}.${spec.name}`);
      }
      assert.equal(observationWidth(manifest.id) > 0, true, manifest.id);
      assert.equal(actionWidth(manifest.id) > 0, true, manifest.id);
    }
  });

  it("refuses a malformed spec rather than encoding it as something", () => {
    assert.throws(() => arenaFieldWidth({ name: "x", kind: "float" }), /finite low < high/);
    assert.throws(
      () => arenaFieldWidth({ name: "x", kind: "float", low: 1, high: 1 }),
      /finite low < high/,
    );
    assert.throws(
      () => arenaFieldWidth({ name: "x", kind: "float", low: 0, high: Number.POSITIVE_INFINITY }),
      /finite low < high/,
    );
    assert.throws(() => arenaFieldWidth({ name: "x", kind: "enum", values: [] }), /non-empty/);
    assert.throws(
      () => arenaFieldWidth({ name: "x", kind: "enum", values: ["a", "a"] }),
      /duplicate/,
    );
  });

  it("throws on an environment id nobody published", () => {
    assert.throws(() => arenaManifest("studio-ops-v2"), /unknown arena environment/);
    assert.throws(() => flattenObservation("nope", {}), /unknown arena environment/);
  });
});

describe("flattenObservation produces a fixed-width vector of finite numbers", () => {
  it("matches the declared width at reset for every environment", () => {
    for (const id of arenaEnvironmentIds()) {
      const environment = ARENA_ENVIRONMENTS[id]!();
      const observation = environment.reset(9, { split: "train" }).observation;
      const vector = flattenObservation(id, observation);
      assert.equal(vector.length, observationWidth(id), id);
      assert.ok(vector.every((value) => Number.isFinite(value)), id);
    }
  });

  it("stays exactly that width at every step of a full 1200-step studio-ops episode", () => {
    // A full episode rather than a short one, because the width has to survive
    // every branch: an empty `nextStationId`, a null payload, a `phase` string
    // the enum has to already know about, and the truncation at the cap.
    const environment = new StudioOpsEnvironment();
    let observation = environment.reset(11, "train-sf-clear-morning-desk-check")
      .observation as StudioOpsObservation;
    const widths = new Set([flattenObservation("studio-ops-v1", observation).length]);
    let steps = 0;
    let terminalReason: string | null = null;
    for (let index = 0; index < 1200; index += 1) {
      // A policy that never works the job but does get the car ready, so the
      // departure is met and the episode runs the whole cap.
      const result = environment.step({
        ...STUDIO_OPS_INACTION,
        micMute: true,
        vehicleCharge: !observation.vehicleReadyByDeparture,
        vehiclePrecondition: Math.abs(observation.vehicleCabinC - 21) > 0.5,
      });
      observation = result.observation as StudioOpsObservation;
      widths.add(flattenObservation("studio-ops-v1", observation).length);
      steps += 1;
      terminalReason = result.info.terminalReason;
      if (result.terminated || result.truncated) break;
    }
    assert.equal(steps, 1200);
    assert.equal(terminalReason, "max-steps");
    assert.deepEqual([...widths], [observationWidth("studio-ops-v1")]);
  });

  it("keeps the width when the observation is empty, partial or malformed", () => {
    const width = observationWidth("studio-ops-v1");
    assert.equal(flattenObservation("studio-ops-v1", {}).length, width);
    assert.equal(flattenObservation("studio-ops-v1", null).length, width);
    assert.equal(flattenObservation("studio-ops-v1", { x: Number.NaN }).length, width);
    assert.ok(
      flattenObservation("studio-ops-v1", { x: Number.NaN }).every((v) => Number.isFinite(v)),
    );
  });

  it("encodes each kind the way the manifest says it does", () => {
    const space = arenaManifest("studio-ops-v1").observationSpace;
    const at = (name: string): number => {
      let cursor = 0;
      for (const spec of space) {
        if (spec.name === name) return cursor;
        cursor += arenaFieldWidth(spec);
      }
      throw new Error(`no field ${name}`);
    };

    // float: the clamped raw value, not a normalization of it.
    const clamped = flattenObservation("studio-ops-v1", { windKph: 5000 });
    assert.equal(clamped[at("windKph")], 120);
    assert.equal(flattenObservation("studio-ops-v1", { windKph: 31.5 })[at("windKph")], 31.5);
    // A field that is absent or unusable falls back to the declared low, which
    // cannot be mistaken for a measurement.
    assert.equal(flattenObservation("studio-ops-v1", {})[at("micLevelDb")], -60);

    // bool: 1 only for a genuine `true`.
    assert.equal(flattenObservation("studio-ops-v1", { deskOccupied: true })[at("deskOccupied")], 1);
    assert.equal(flattenObservation("studio-ops-v1", { deskOccupied: 1 })[at("deskOccupied")], 0);

    // enum: one-hot, and all-zeros for a value outside the vocabulary — which
    // is what a null payload has to encode as. "Carrying nothing" is not a
    // thing being carried, and must not collide with "carrying a parcel".
    const cursor = at("weatherCondition");
    const rain = flattenObservation("studio-ops-v1", { weatherCondition: "rain" });
    assert.equal(rain.slice(cursor, cursor + 8).reduce((sum, value) => sum + value, 0), 1);
    assert.equal(rain[cursor + 5], 1);
    const unknown = flattenObservation("studio-ops-v1", { weatherCondition: "hail" });
    assert.deepEqual(unknown.slice(cursor, cursor + 8), [0, 0, 0, 0, 0, 0, 0, 0]);
    assert.equal(flattenObservation("studio-ops-v1", { payload: null })[at("payload")], 0);
    assert.equal(flattenObservation("studio-ops-v1", { payload: "parcel" })[at("payload")], 1);

    // id: one stable slot in [0, 1), changing exactly when the identity does.
    const first = flattenObservation("studio-ops-v1", { nextStationId: "sf-studio-monitor" });
    const same = flattenObservation("studio-ops-v1", { nextStationId: "sf-studio-monitor" });
    const other = flattenObservation("studio-ops-v1", { nextStationId: "sf-studio-display" });
    assert.equal(first[at("nextStationId")], same[at("nextStationId")]);
    assert.notEqual(first[at("nextStationId")], other[at("nextStationId")]);
    assert.ok(first[at("nextStationId")]! >= 0 && first[at("nextStationId")]! < 1);
    assert.equal(flattenObservation("studio-ops-v1", { nextStationId: "" })[at("nextStationId")], 0);
  });

  it("never observes an enum value its own vocabulary does not carry", () => {
    // The one way a one-hot silently loses information: a `phase` string the
    // controller assigns and the manifest has never heard of encodes as all
    // zeros and reads to a policy as an unremarkable state.
    const space = arenaManifest("studio-ops-v1").observationSpace;
    const enums = space.filter((spec) => spec.kind === "enum");
    const seen = new Map<string, Set<string>>(enums.map((spec) => [spec.name, new Set<string>()]));
    const record = (observation: StudioOpsObservation): void => {
      for (const spec of enums) {
        const value = (observation as unknown as Record<string, unknown>)[spec.name];
        if (typeof value === "string") seen.get(spec.name)!.add(value);
      }
    };
    // Three policies, because one policy visits one corridor of the phase
    // machine: the scripted baseline works its stations, the drifter never
    // interacts and sits in `awaiting-interaction`, and the spammer walks into
    // walls and reaches the recovery phases.
    const POLICIES: ((observation: StudioOpsObservation, step: number) => StudioOpsAction)[] = [
      (observation) => studioOpsScriptedBaseline(observation),
      (observation, step) => ({
        ...STUDIO_OPS_INACTION,
        micMute: true,
        x: Math.sin(step / 40) * 0.4,
        z: Math.cos(step / 37) * 0.4,
        speakerPlay: observation.deskOccupied,
      }),
      () => ({ ...STUDIO_OPS_INACTION, z: 1, interact: true }),
    ];
    for (const definition of STUDIO_OPS_SCENARIOS.definitions) {
      for (const seed of [4, 77]) {
        for (const policy of POLICIES) {
          const environment = new StudioOpsEnvironment();
          let observation = environment.reset(seed, definition.id).observation as
            StudioOpsObservation;
          record(observation);
          for (let index = 0; index < 900; index += 1) {
            const result = environment.step(policy(observation, index));
            observation = result.observation as StudioOpsObservation;
            record(observation);
            if (result.terminated || result.truncated) break;
          }
        }
      }
    }
    for (const spec of enums) {
      for (const value of seen.get(spec.name)!) {
        assert.ok(spec.values!.includes(value), `${spec.name} observed unknown "${value}"`);
      }
    }
    // And the sweep must actually have seen something, or this proves nothing.
    assert.ok(seen.get("phase")!.size >= 6, [...seen.get("phase")!].join(","));
    assert.ok(seen.get("mode")!.size >= 3, [...seen.get("mode")!].join(","));
    assert.ok(seen.get("weatherCondition")!.size >= 3, [...seen.get("weatherCondition")!].join(","));
    assert.ok(seen.get("levelId")!.size === 2, [...seen.get("levelId")!].join(","));
  });
});

describe("structureAction inverts the action encoding", () => {
  it("round-trips every environment's own baseline action", () => {
    for (const id of arenaEnvironmentIds()) {
      const environment = ARENA_ENVIRONMENTS[id]!();
      environment.reset(2, { split: "train" });
      const zeroed = structureAction(id, new Array(actionWidth(id)).fill(0));
      const vector = flattenAction(id, zeroed);
      assert.equal(vector.length, actionWidth(id), id);
      assert.deepEqual(structureAction(id, vector), zeroed, id);
      // And the result is a legal action: stepping with it must not throw.
      assert.ok(Number.isFinite(environment.step(zeroed).reward), id);
    }
  });

  it("clamps a raw network output instead of refusing it", () => {
    const action = structureAction("studio-ops-v1", [
      9, -9, 1, 999, 0.4, 5, -1, 0.5, 0.49,
    ]);
    assert.equal(action.x, 1);
    assert.equal(action.z, -1);
    assert.equal(action.interact, true);
    assert.equal(action.micGain, 36);
    assert.equal(action.micMute, false);
    assert.equal(action.speakerVolume, 1);
    assert.equal(action.speakerPlay, false);
    assert.equal(action.vehiclePrecondition, true);
    assert.equal(action.vehicleCharge, false);
  });

  it("refuses a vector of the wrong length", () => {
    assert.throws(() => structureAction("studio-ops-v1", [0, 0]), /exactly 9 values/);
    assert.throws(
      () => structureAction("studio-ops-v1", new Array(10).fill(0)),
      /exactly 9 values/,
    );
  });
});
