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

// --- pure helpers ------------------------------------------------------------------------------

function inst(partial: Partial<InstanceRecord> & { instanceId: string }): InstanceRecord {
  return {
    state: "ready",
    sessions: 0,
    epoch: 1,
    lastSeenAt: NOW,
    lastActiveAt: NOW,
    tier: partial.instanceId === "host" ? 0 : 1,
    ...partial,
  };
}

const NOW = 1_000_000_000_000;
const SETTINGS: PoolSettings = { ...DEFAULT_SETTINGS };

test("clampSettings bounds maxContainers and idle timeout, keeps unknown-free", () => {
  const s = clampSettings(DEFAULT_SETTINGS, { maxContainers: 99, idleTimeoutSeconds: 5 });
  assert.equal(s.maxContainers, 16);
  assert.equal(s.idleTimeoutSeconds, 60);
  const s2 = clampSettings(DEFAULT_SETTINGS, { maxContainers: 0 });
  assert.equal(s2.maxContainers, 1);
  const s3 = clampSettings(DEFAULT_SETTINGS, { balancerPolicy: "bogus" as never });
  assert.equal(s3.balancerPolicy, DEFAULT_SETTINGS.balancerPolicy);
  const s4 = clampSettings(DEFAULT_SETTINGS, null);
  assert.deepEqual(s4, DEFAULT_SETTINGS);
});

test("least-sessions places on the emptiest live sandbox, ties break deterministically", () => {
  const placement = placeSandbox({
    instances: [
      inst({ instanceId: "host", sessions: 3 }),
      inst({ instanceId: "sb-1", sessions: 2 }),
      inst({ instanceId: "sb-2", sessions: 1 }),
      inst({ instanceId: "sb-3", sessions: 1 }),
    ],
    settings: SETTINGS,
    rrCursor: 0,
    now: NOW,
  });
  assert.deepEqual(placement, { decision: "place", instanceId: "sb-2" });
});

test("round-robin walks the live sandboxes by cursor", () => {
  const instances = [inst({ instanceId: "sb-1" }), inst({ instanceId: "sb-2" }), inst({ instanceId: "sb-3" })];
  const settings: PoolSettings = { ...SETTINGS, balancerPolicy: "round-robin" };
  const picks = [0, 1, 2, 3].map(
    (rrCursor) => placeSandbox({ instances, settings, rrCursor, now: NOW }) as { instanceId: string },
  );
  assert.deepEqual(
    picks.map((p) => p.instanceId),
    ["sb-1", "sb-2", "sb-3", "sb-1"],
  );
});

test("stale or starting sandboxes are not placeable", () => {
  const placement = placeSandbox({
    instances: [
      inst({ instanceId: "host" }),
      inst({ instanceId: "sb-1", lastSeenAt: NOW - HEARTBEAT_STALE_MS - 1 }),
      inst({ instanceId: "sb-2", state: "starting" }),
    ],
    settings: { ...SETTINGS, maxContainers: 3 },
    rrCursor: 0,
    now: NOW,
  });
  // Both sandboxes exist (cap reached at 3 total) but neither is placeable → at capacity.
  assert.deepEqual(placement, { decision: "at_capacity" });
});

test("boots a new sandbox under the cap, host counts against it", () => {
  const under = placeSandbox({
    instances: [inst({ instanceId: "host" })],
    settings: { ...SETTINGS, maxContainers: 2 },
    rrCursor: 0,
    now: NOW,
  });
  assert.deepEqual(under, { decision: "boot", instanceId: "sb-1" });

  const at = placeSandbox({
    instances: [inst({ instanceId: "host" })],
    settings: { ...SETTINGS, maxContainers: 1 },
    rrCursor: 0,
    now: NOW,
  });
  assert.deepEqual(at, { decision: "at_capacity" });
});

test("boot picks the first free sb-<n> slot (stopped slots are reusable)", () => {
  const placement = placeSandbox({
    instances: [
      inst({ instanceId: "host" }),
      inst({ instanceId: "sb-1", state: "starting" }), // occupied but not placeable
      inst({ instanceId: "sb-2", state: "stopped" }), // free slot
    ],
    settings: { ...SETTINGS, maxContainers: 4 },
    rrCursor: 0,
    now: NOW,
  });
  assert.deepEqual(placement, { decision: "boot", instanceId: "sb-2" });
});

// --- 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,
  });
}

test("provision boots the host container once and is idempotent", async () => {
  const { core, started } = makeCore();
  const first = (await (await core.fetch(req("POST", "/provision", "pool"))).json()) as Record<string, unknown>;
  assert.equal(first.instanceId, "host");
  assert.equal(first.alreadyRunning, false);
  assert.equal(started.length, 1);
  assert.equal(started[0].env.NUCLEIC_RUNNER_INSTANCE_TOKEN, "tok-pool-1-host");

  const second = (await (await core.fetch(req("POST", "/provision", "pool"))).json()) as Record<string, unknown>;
  assert.equal(second.alreadyRunning, true);
  assert.equal(started.length, 1);
});

test("acquire boots a sandbox, then places on it once it heartbeats ready", async () => {
  let now = NOW;
  const { core, started } = makeCore({ now: () => now });
  await core.fetch(req("POST", "/provision", "pool"));

  const boot = (await (await core.fetch(req("POST", "/sandbox/acquire", "instance:host"))).json()) as Record<
    string,
    unknown
  >;
  assert.equal(boot.decision, "boot");
  assert.equal(boot.instanceId, "sb-1");
  assert.equal(started.length, 2);

  // The new sandbox heartbeats in with its boot epoch → becomes placeable.
  const hb = await core.fetch(req("POST", "/heartbeat", "instance:sb-1", { epoch: boot.epoch, sessions: 1 }));
  assert.equal(hb.status, 200);
  const minted = (await hb.json()) as { token: string };
  assert.equal(minted.token, "tok-pool-1-sb-1");

  const place = (await (await core.fetch(req("POST", "/sandbox/acquire", "instance:host"))).json()) as Record<
    string,
    unknown
  >;
  assert.deepEqual(
    { decision: place.decision, instanceId: place.instanceId },
    { decision: "place", instanceId: "sb-1" },
  );
});

test("acquire reports at_capacity at the cap instead of failing", async () => {
  const { core } = makeCore();
  await core.fetch(req("PATCH", "/settings", "pool", { maxContainers: 1 }));
  await core.fetch(req("POST", "/provision", "pool"));
  const res = (await (await core.fetch(req("POST", "/sandbox/acquire", "instance:host"))).json()) as Record<
    string,
    unknown
  >;
  assert.equal(res.decision, "at_capacity");
});

test("heartbeat with a stale epoch is fenced with 409", async () => {
  const { core } = makeCore();
  await core.fetch(req("POST", "/provision", "pool"));
  const boot = (await (await core.fetch(req("POST", "/sandbox/acquire", "instance:host"))).json()) as {
    epoch: number;
  };
  const stale = await core.fetch(req("POST", "/heartbeat", "instance:sb-1", { epoch: boot.epoch - 1, sessions: 0 }));
  assert.equal(stale.status, 409);
  const fresh = await core.fetch(req("POST", "/heartbeat", "instance:sb-1", { epoch: boot.epoch, sessions: 0 }));
  assert.equal(fresh.status, 200);
});

test("release decrements sessions; alarm reaps idle sandboxes past the 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"));
  const boot = (await (await core.fetch(req("POST", "/sandbox/acquire", "instance:host"))).json()) as {
    epoch: number;
  };
  await core.fetch(req("POST", "/heartbeat", "instance:sb-1", { epoch: boot.epoch, sessions: 1 }));
  await core.fetch(req("POST", "/sandbox/release", "instance:host", { instanceId: "sb-1" }));

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

  now += 61_000;
  await core.alarm();
  assert.deepEqual(stopped, ["sb-1"]);

  const status = (await (await core.fetch(req("GET", "/status", "pool"))).json()) as {
    instances: Array<{ instanceId: string; state: string }>;
  };
  assert.equal(status.instances.find((i) => i.instanceId === "sb-1")?.state, "stopped");
});

test("the host container is never reaped by the idle alarm", 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"));
  now += 10 * 60_000;
  await core.alarm();
  assert.deepEqual(stopped, []);
});

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

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

test("management verbs require pool auth; heartbeat requires instance auth", async () => {
  const { core } = makeCore();
  assert.equal((await core.fetch(req("GET", "/status", "instance:host"))).status, 403);
  assert.equal((await core.fetch(req("PATCH", "/settings", "instance:host", {}))).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);
});

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