import assert from "node:assert/strict";
import { test } from "node:test";
import { StatsCache, type FetchFn } from "../src/lib/stats-cache";
import { type LlamaSwapConfig } from "../src/lib/util";

const cfg: LlamaSwapConfig = { baseUrl: "http://x" };

function stubFetch(result: unknown, counter: { count: number }): FetchFn {
  return async () => {
    counter.count++;
    return result as never;
  };
}

async function flush(): Promise<void> {
  await new Promise((r) => setTimeout(r, 0));
}

test("register polls once immediately and caches the value", async () => {
  const calls = { count: 0 };
  const cache = new StatsCache(cfg, stubFetch({ totalRequests: 3, totalInputTokens: 1, totalOutputTokens: 2, genP95: 4 }, calls), 1000, 30);
  cache.register("all");
  await flush();
  assert.equal(calls.count, 1);
  assert.equal(cache.get("all")!.totalRequests, 3);
  cache.unregister("all");
});

test("multiple keys are each polled", async () => {
  const keys: string[] = [];
  const cache = new StatsCache(
    cfg,
    async (_c, key) => {
      keys.push(key);
      return { totalRequests: 1, totalInputTokens: 1, totalOutputTokens: 1, genP95: 1 };
    },
    1000,
    30,
  );
  cache.register("a");
  cache.register("b");
  await flush();
  assert.deepEqual(keys.sort(), ["a", "b"]);
  cache.unregister("a");
  cache.unregister("b");
});

test("unregister stops the interval and clears the value", async () => {
  const calls = { count: 0 };
  const cache = new StatsCache(cfg, stubFetch({ totalRequests: 1, totalInputTokens: 1, totalOutputTokens: 1, genP95: 1 }, calls), 10, 30);
  cache.register("a");
  await flush();
  assert.equal(calls.count, 1);
  cache.unregister("a");
  await new Promise((r) => setTimeout(r, 40));
  assert.equal(cache.get("a"), undefined);
  assert.equal(calls.count, 1);
});

test("scheduleRefresh throttles to one poll per window", async () => {
  const calls = { count: 0 };
  const cache = new StatsCache(cfg, stubFetch({ totalRequests: 1, totalInputTokens: 1, totalOutputTokens: 1, genP95: 1 }, calls), 10000, 30);
  cache.register("a");
  await flush();
  const before = calls.count;
  cache.scheduleRefresh();
  cache.scheduleRefresh();
  cache.scheduleRefresh();
  await flush();
  assert.equal(calls.count, before);
  await new Promise((r) => setTimeout(r, 60));
  assert.equal(calls.count, before + 1);
  cache.unregister("a");
});

test("fetch failure keeps the last-known value", async () => {
  const fail: FetchFn = async () => {
    throw new Error("boom");
  };
  const cache = new StatsCache(cfg, fail, 10000, 30);
  cache.register("a");
  await flush();
  assert.equal(cache.get("a"), undefined);
  cache.unregister("a");
});

test("ref-counted register: a sibling unregister keeps the key active", async () => {
  const calls = { count: 0 };
  const cache = new StatsCache(cfg, stubFetch({ totalRequests: 5, totalInputTokens: 1, totalOutputTokens: 1, genP95: 1 }, calls), 10000, 30);
  cache.register("all");
  await flush();
  assert.equal(cache.get("all")!.totalRequests, 5);
  cache.register("all");
  cache.unregister("all");
  await flush();
  assert.equal(cache.get("all")!.totalRequests, 5);
  cache.unregister("all");
  assert.equal(cache.get("all"), undefined);
});

test("setConfig keeps registrations but clears cached values", async () => {
  const cache = new StatsCache(cfg, async (_c, key) => ({ totalRequests: 5, totalInputTokens: 1, totalOutputTokens: 1, genP95: 1 }), 10000, 30);
  cache.register("all");
  await flush();
  assert.equal(cache.get("all")!.totalRequests, 5);
  cache.setConfig({ baseUrl: "http://new" });
  assert.equal(cache.get("all"), undefined);
  cache.scheduleRefresh();
  await flush();
  await new Promise((r) => setTimeout(r, 60));
  assert.equal(cache.get("all")!.totalRequests, 5);
  cache.unregister("all");
});