import { test } from "node:test";
import assert from "node:assert/strict";
import {
  BOOT_GRACE_MS,
  clampSettings,
  DEFAULT_SETTINGS,
  DONE_SLEEP_MS,
  HEARTBEAT_STALE_MS,
  RunnerPoolCore,
  type ContainerDriver,
  type PoolSettings,
} from "../src/pool.ts";

const NOW = 1_000_000_000_000;

// --- clampSettings (pure) ----------------------------------------------------------------------

test("clampSettings bounds maxRunners and idle timeout, drops the removed balancer knob", () => {
  const s = clampSettings(DEFAULT_SETTINGS, { maxRunners: 99, idleTimeoutSeconds: 5 });
  assert.equal(s.maxRunners, 16);
  assert.equal(s.idleTimeoutSeconds, 60);
  const s2 = clampSettings(DEFAULT_SETTINGS, { maxRunners: 0 });
  assert.equal(s2.maxRunners, 1);
  // The old balancerPolicy field no longer exists — an unknown key is simply ignored.
  const s3 = clampSettings(DEFAULT_SETTINGS, { balancerPolicy: "round-robin" } as never);
  assert.deepEqual(s3, DEFAULT_SETTINGS);
  const s4 = clampSettings(DEFAULT_SETTINGS, null);
  assert.deepEqual(s4, DEFAULT_SETTINGS);
});

test("clampSettings accepts only known intelligence modes (item 6)", () => {
  assert.equal(DEFAULT_SETTINGS.intelligenceMode, "mesh");
  const agent = clampSettings(DEFAULT_SETTINGS, { intelligenceMode: "agent" });
  assert.equal(agent.intelligenceMode, "agent");
  const heuristic = clampSettings(agent, { intelligenceMode: "heuristic" });
  assert.equal(heuristic.intelligenceMode, "heuristic");
  // Garbage keeps the current value — the DO never trusts a client.
  const bogus = clampSettings(agent, { intelligenceMode: "gpt-99" as never });
  assert.equal(bogus.intelligenceMode, "agent");
  // A patch that doesn't mention it leaves it untouched.
  const untouched = clampSettings(agent, { maxRunners: 2 });
  assert.equal(untouched.intelligenceMode, "agent");
});

// --- RunnerPoolCore over fakes -----------------------------------------------------------------

function fakeStorage() {
  const map = new Map<string, unknown>();
  let alarm: number | null = null;
  return {
    map,
    get alarmTime() {
      return alarm;
    },
    async get<T>(key: string) {
      return map.get(key) as T | undefined;
    },
    async put(key: string, value: unknown) {
      map.set(key, value);
    },
    async delete(key: string) {
      return map.delete(key);
    },
    async setAlarm(time: number) {
      alarm = time;
    },
  };
}

function fakeDriver() {
  const started: Array<{ instanceId: string; env: Record<string, string> }> = [];
  const stopped: string[] = [];
  const driver: ContainerDriver = {
    async start(_pool, instanceId, env) {
      started.push({ instanceId, env });
    },
    async stop(_pool, instanceId) {
      stopped.push(instanceId);
    },
  };
  return { driver, started, stopped };
}

function makeCore(opts?: { now?: () => number }) {
  const storage = fakeStorage();
  const { driver, started, stopped } = fakeDriver();
  const core = new RunnerPoolCore({ storage }, "pool-1", {
    driver,
    mintInstanceToken: async (pool, instance) => `tok-${pool}-${instance}`,
    now: opts?.now,
  });
  return { core, storage, started, stopped };
}

function req(method: string, path: string, auth: string, body?: unknown): Request {
  return new Request(`https://pool${path}`, {
    method,
    headers: { "x-nucleic-auth": auth, ...(body !== undefined ? { "content-type": "application/json" } : {}) },
    body: body !== undefined ? JSON.stringify(body) : undefined,
  });
}

async function bodyOf(res: Response): Promise<Record<string, unknown>> {
  return (await res.json()) as Record<string, unknown>;
}

// --- provision (demand-driven scale-up) --------------------------------------------------------

test("provision boots the first runner slot, then grows the fleet one runner at a time", async () => {
  const { core, started } = makeCore();
  const first = await bodyOf(await core.fetch(req("POST", "/provision", "pool")));
  assert.equal(first.decision, "boot");
  assert.equal(first.instanceId, "runner-1");
  assert.equal(started.length, 1);
  assert.equal(started[0].env.NUCLEIC_RUNNER_INSTANCE_ID, "runner-1");
  assert.equal(started[0].env.NUCLEIC_RUNNER_INSTANCE_TOKEN, "tok-pool-1-runner-1");
  // Intelligence mode rides every boot (item 6) — the default when never patched.
  assert.equal(started[0].env.NUCLEIC_RUNNER_INTELLIGENCE_MODE, "mesh");

  const second = await bodyOf(await core.fetch(req("POST", "/provision", "pool")));
  assert.equal(second.decision, "boot");
  assert.equal(second.instanceId, "runner-2");
  assert.equal(started.length, 2);
});

test("provision reports at_capacity at the maxRunners cap instead of failing", async () => {
  const { core, started } = makeCore();
  await core.fetch(req("PATCH", "/settings", "pool", { maxRunners: 2 }));
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/provision", "pool"));
  const capped = await bodyOf(await core.fetch(req("POST", "/provision", "pool")));
  assert.equal(capped.decision, "at_capacity");
  assert.equal(capped.liveCount, 2);
  assert.equal(capped.maxRunners, 2);
  assert.equal(started.length, 2);
});

test("provision reuses the lowest free slot after one is reaped/stopped", async () => {
  const { core, started } = makeCore();
  await core.fetch(req("POST", "/provision", "pool")); // runner-1
  await core.fetch(req("POST", "/provision", "pool")); // runner-2
  await core.fetch(req("POST", "/stop", "pool")); // both → stopped, slots free
  const next = await bodyOf(await core.fetch(req("POST", "/provision", "pool")));
  assert.equal(next.instanceId, "runner-1"); // lowest free slot reused
  assert.equal(started.length, 3);
});

test("a patched intelligence mode reaches the next runner boot (item 6)", async () => {
  const { core, started } = makeCore();
  const patched = (await bodyOf(
    await core.fetch(req("PATCH", "/settings", "pool", { intelligenceMode: "agent" })),
  )) as unknown as { settings: PoolSettings };
  assert.equal(patched.settings.intelligenceMode, "agent");
  await core.fetch(req("POST", "/provision", "pool"));
  assert.equal(started[0].env.NUCLEIC_RUNNER_INTELLIGENCE_MODE, "agent");
});

test("settings stored before the field rename read defaults and drop legacy keys", async () => {
  const { core, storage } = makeCore();
  // A pre-rename pool persisted maxContainers + balancerPolicy.
  await storage.put("settings", { maxContainers: 2, balancerPolicy: "round-robin", idleTimeoutSeconds: 120 });
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    settings: Record<string, unknown>;
  };
  assert.equal(status.settings.intelligenceMode, "mesh");
  assert.equal(status.settings.maxRunners, DEFAULT_SETTINGS.maxRunners); // legacy maxContainers ignored
  assert.equal(status.settings.idleTimeoutSeconds, 120); // a still-valid field survives
  assert.equal("balancerPolicy" in status.settings, false); // dropped
  assert.equal("maxContainers" in status.settings, false); // dropped
});

// --- heartbeat / fencing -----------------------------------------------------------------------

test("heartbeat re-mints the instance token and fences a stale epoch with 409", async () => {
  const { core } = makeCore();
  const boot = await bodyOf(await core.fetch(req("POST", "/provision", "pool")));
  const epoch = boot.epoch as number;
  const hb = await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch, sessions: 1 }));
  assert.equal(hb.status, 200);
  assert.equal((await bodyOf(hb)).token, "tok-pool-1-runner-1");

  const stale = await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: epoch - 1, sessions: 0 }));
  assert.equal(stale.status, 409);
});

// --- wake (relay presence → a live slot is a no-op; anything else is declined) ------------------

test("a targeted wake never reboots — live is alreadyRunning, dead/stale is declined", async () => {
  let now = NOW;
  const { core, started } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool")); // runner-1
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 }));

  // Still fresh → no reboot.
  const live = await bodyOf(await core.fetch(req("POST", "/wake", "pool", { instanceId: "runner-1" })));
  assert.equal(live.alreadyRunning, true);
  assert.equal(started.length, 1);

  // Gone stale → DECLINED, not rebooted: a reboot mints a fresh in-container identity that can
  // never serve the room the device is dialing — it would only boot a mesh-orphan.
  now += HEARTBEAT_STALE_MS + 1;
  const declined = await bodyOf(await core.fetch(req("POST", "/wake", "pool", { instanceId: "runner-1" })));
  assert.equal(declined.decision, "identityRotated");
  assert.equal(started.length, 1); // nothing booted

  // No epoch bump happened, so a wedged-but-alive container that resumes heartbeating recovers.
  const recovered = await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 }));
  assert.equal(recovered.status, 200);

  // A reaped (stopped) slot is likewise declined rather than rebooted into an orphan.
  await core.fetch(req("POST", "/stop", "pool"));
  const stopped = await bodyOf(await core.fetch(req("POST", "/wake", "pool", { instanceId: "runner-1" })));
  assert.equal(stopped.decision, "identityRotated");
  assert.equal(started.length, 1);
});

test("a bodyless wake behaves like provision (ensure one runner under the cap)", async () => {
  const { core, started } = makeCore();
  const woken = await bodyOf(await core.fetch(req("POST", "/wake", "pool")));
  assert.equal(woken.decision, "boot");
  assert.equal(woken.instanceId, "runner-1");
  assert.equal(started.length, 1);
});

// --- alarm: idle reap, dead reap, boot grace, scale-to-cap -------------------------------------

test("alarm scales an idle runner to zero past the idle timeout", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("PATCH", "/settings", "pool", { idleTimeoutSeconds: 60 }));
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 }));

  await core.alarm(); // not idle long enough
  assert.deepEqual(stopped, []);

  now += 61_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 })); // fresh, still idle
  await core.alarm();
  assert.deepEqual(stopped, ["runner-1"]);
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ instanceId: string; state: string }>;
  };
  assert.equal(status.instances.find((i) => i.instanceId === "runner-1")?.state, "stopped");
});

test("an idle runner stays up while it holds a session", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("PATCH", "/settings", "pool", { idleTimeoutSeconds: 60 }));
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1 }));
  now += 61_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1 }));
  await core.alarm();
  assert.deepEqual(stopped, []);
});

test("alarm reaps a dead runner even when it last reported live sessions", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 3 }));

  await core.alarm(); // fresh + sessions ⇒ untouched
  assert.deepEqual(stopped, []);

  // Heartbeats stop but the ghost record still says sessions: 3 — without the dead-reap it would
  // hold a maxRunners slot forever.
  now += HEARTBEAT_STALE_MS + 1;
  await core.alarm();
  assert.deepEqual(stopped, ["runner-1"]);
});

test("a still-starting runner gets the boot grace before being declared a failed boot", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool")); // runner-1 stays "starting" — never heartbeats

  // Past the ready-staleness window but within the boot grace: a slow cold pull, not dead.
  now += HEARTBEAT_STALE_MS + 1;
  await core.alarm();
  assert.deepEqual(stopped, []);

  // Past the boot grace: it never came up — reap the wedged slot.
  now = NOW + BOOT_GRACE_MS + 1;
  await core.alarm();
  assert.deepEqual(stopped, ["runner-1"]);
});

test("a lowered cap drains idle runners down to maxRunners (oldest-first)", async () => {
  const { core, stopped } = makeCore({ now: () => NOW });
  // Three fresh, idle-but-not-timed-out runners.
  await core.fetch(req("POST", "/provision", "pool")); // runner-1
  await core.fetch(req("POST", "/provision", "pool")); // runner-2
  await core.fetch(req("POST", "/provision", "pool")); // runner-3
  for (const n of [1, 2, 3]) {
    await core.fetch(req("POST", "/heartbeat", `instance:runner-${n}`, { epoch: n, sessions: 0 }));
  }
  await core.fetch(req("PATCH", "/settings", "pool", { maxRunners: 1 }));
  await core.alarm();
  // Two runners trimmed to reach the cap of 1; one survives.
  assert.equal(stopped.length, 2);
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ instanceId: string; state: string }>;
  };
  assert.equal(status.instances.filter((i) => i.state !== "stopped").length, 1);
});

// --- alarm: all-done sleep (sessions remain, but every one is Done and quiet) -------------------

test("alarm sleeps a runner once all its sessions are Done and settled past the window", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));
  // Two live sessions, both actively working — nowhere near sleep.
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 2, activeSessions: 2 }));
  await core.alarm();
  assert.deepEqual(stopped, []);

  // Both turns finish: sessions still present (the finished chats sit on the runner), but none is
  // active. The done-sleep clock starts here.
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 2, activeSessions: 0 }));
  await core.alarm(); // just went quiet — not settled long enough
  assert.deepEqual(stopped, []);

  // Still quiet, but not yet past DONE_SLEEP_MS.
  now += DONE_SLEEP_MS - 1_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 2, activeSessions: 0 }));
  await core.alarm();
  assert.deepEqual(stopped, []);

  // Past the window with no active session the whole time → sleep it.
  now += 2_000;
  await core.alarm();
  assert.deepEqual(stopped, ["runner-1"]);
});

test("a queued follow-up on a Done session keeps the runner awake (activeSessions > 0)", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));
  // One session is Done but still holds a queued message waiting to dispatch — the runner reports
  // it as active, so the pool must not sleep even long past the window.
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1, activeSessions: 1 }));
  now += DONE_SLEEP_MS + 60_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1, activeSessions: 1 }));
  await core.alarm();
  assert.deepEqual(stopped, []);
});

test("a fresh active turn resets the done-sleep clock (activeSessions goes back up)", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1, activeSessions: 0 }));

  // Most of the way to sleep, then the user sends a follow-up: the session is active again.
  now += DONE_SLEEP_MS - 1_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1, activeSessions: 1 }));
  now += 5_000; // past the original deadline, but the clock reset when it went active
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 1, activeSessions: 0 }));
  await core.alarm();
  assert.deepEqual(stopped, []);

  // Now let it settle for the full window from the reset.
  now += DONE_SLEEP_MS + 1_000;
  await core.alarm();
  assert.deepEqual(stopped, ["runner-1"]);
});

test("a runner from an older Worker (no activeSessions) never sleeps down the all-done path", async () => {
  let now = NOW;
  const { core, stopped } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));
  // Heartbeats omit activeSessions (pre-field runner) — the pool assumes all sessions are active
  // and leaves the runner up no matter how long it holds them.
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 2 }));
  now += DONE_SLEEP_MS + 60_000;
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 2 }));
  await core.alarm();
  assert.deepEqual(stopped, []);
});

test("status surfaces the runner's active-session count", async () => {
  const { core } = makeCore({ now: () => NOW });
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 3, activeSessions: 1 }));
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ sessions: number; activeSessions: number }>;
  };
  assert.equal(status.instances[0].sessions, 3);
  assert.equal(status.instances[0].activeSessions, 1);
});

// --- status / stop / auth ----------------------------------------------------------------------

test("status clamps maxRunners server-side and reports staleness", async () => {
  let now = NOW;
  const { core } = makeCore({ now: () => now });
  const patched = (await bodyOf(await core.fetch(req("PATCH", "/settings", "pool", { maxRunners: 500 })))) as unknown as {
    settings: PoolSettings;
  };
  assert.equal(patched.settings.maxRunners, 16);

  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 })); // ready
  now += HEARTBEAT_STALE_MS + 1;
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ instanceId: string; stale: boolean }>;
  };
  assert.equal(status.instances[0].stale, true);
});

test("stop stops every live runner and zeroes sessions", async () => {
  const { core, stopped } = makeCore();
  await core.fetch(req("POST", "/provision", "pool")); // runner-1
  await core.fetch(req("POST", "/provision", "pool")); // runner-2
  await core.fetch(req("POST", "/heartbeat", "instance:runner-2", { epoch: 2, sessions: 2 }));
  await core.fetch(req("POST", "/stop", "pool"));
  assert.deepEqual(stopped.sort(), ["runner-1", "runner-2"]);
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ state: string; sessions: number }>;
  };
  for (const i of status.instances) {
    assert.equal(i.state, "stopped");
    assert.equal(i.sessions, 0);
  }
});

test("heartbeat records the runner's mesh roomId; a reboot clears it (fresh identity)", async () => {
  let now = NOW;
  const { core } = makeCore({ now: () => now });
  const room = "ab".repeat(32);
  await core.fetch(req("POST", "/provision", "pool"));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0, roomId: room }));
  let status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ roomId?: string }>;
  };
  assert.equal(status.instances[0].roomId, room);

  // A malformed roomId is dropped; the recorded one survives a roomId-less heartbeat.
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0, roomId: "nope" }));
  await core.fetch(req("POST", "/heartbeat", "instance:runner-1", { epoch: 1, sessions: 0 }));
  status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ roomId?: string }>;
  };
  assert.equal(status.instances[0].roomId, room);

  // A reboot of the slot mints a fresh in-container identity — the old roomId must not linger
  // (a device would mistake the orphan for already-joined).
  now += HEARTBEAT_STALE_MS + 1;
  await core.fetch(req("POST", "/stop", "pool"));
  await core.fetch(req("POST", "/provision", "pool")); // runner-1 reborn on the freed slot
  status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ roomId?: string }>;
  };
  assert.equal(status.instances[0].roomId, undefined);
});

test("prune deletes stopped rows only, freeing the list without touching live runners", async () => {
  const { core } = makeCore();
  await core.fetch(req("POST", "/provision", "pool")); // runner-1
  await core.fetch(req("POST", "/provision", "pool")); // runner-2
  await core.fetch(req("POST", "/stop", "pool")); // both stopped
  await core.fetch(req("POST", "/provision", "pool")); // runner-1 reborn (live)

  const pruned = await bodyOf(await core.fetch(req("POST", "/prune", "pool")));
  assert.equal(pruned.removed, 1); // only the stopped runner-2 row
  const status = (await bodyOf(await core.fetch(req("GET", "/status", "pool")))) as unknown as {
    instances: Array<{ instanceId: string; state: string }>;
  };
  assert.deepEqual(status.instances.map((i) => i.instanceId), ["runner-1"]);
  assert.notEqual(status.instances[0].state, "stopped");

  // Idempotent — nothing left to remove.
  const again = await bodyOf(await core.fetch(req("POST", "/prune", "pool")));
  assert.equal(again.removed, 0);
});

test("management verbs require pool auth; heartbeat requires instance auth", async () => {
  const { core } = makeCore();
  assert.equal((await core.fetch(req("GET", "/status", "instance:runner-1"))).status, 403);
  assert.equal((await core.fetch(req("PATCH", "/settings", "instance:runner-1", {}))).status, 403);
  assert.equal((await core.fetch(req("POST", "/provision", "instance:runner-1"))).status, 403);
  assert.equal((await core.fetch(req("POST", "/prune", "instance:runner-1"))).status, 403);
  assert.equal((await core.fetch(req("POST", "/heartbeat", "pool", { epoch: 1 }))).status, 403);
  assert.equal((await core.fetch(req("POST", "/heartbeat", "", { epoch: 1 }))).status, 403);
});
