import assert from "node:assert/strict";
import test from "node:test";

// Node's type-stripping test runner requires the explicit TypeScript suffix.
// @ts-expect-error The application compiler intentionally uses extensionless imports.
import {
  GENERATION_IDLE_MS,
  ROLLING_WINDOW_MS,
  charsToTokensPerSec,
  messageHasToolCalls,
  rollingTokensPerSec,
  shouldFreezeGeneration,
  usageTokensPerSec,
} from "./speed-metrics.ts";

test("rolling window reports the recent rate while characters grow", () => {
  const samples = [
    { at: 0, chars: 0 },
    { at: 500, chars: 100 },
    { at: 1000, chars: 200 },
  ];
  const rate = rollingTokensPerSec(samples, 200, 1000, 1000);
  // 200 chars / 1000ms / 2 chars-per-token = 100 tok/s
  assert.equal(rate, 100);
});

test("rolling window drops to 0 when every sample is older than the window", () => {
  const samples = [
    { at: 0, chars: 0 },
    { at: 200, chars: 80 },
    { at: 400, chars: 160 },
  ];
  const now = 400 + ROLLING_WINDOW_MS + 50;
  assert.equal(rollingTokensPerSec(samples, 160, now), 0);
});

test("rolling window decays toward 0 after a mid-stream pause", () => {
  const samples = [
    { at: 0, chars: 0 },
    { at: 400, chars: 80 },
  ];
  const live = rollingTokensPerSec(samples, 80, 500, 1000);
  const paused = rollingTokensPerSec(samples, 80, 900, 1000);
  assert.ok(live > paused);
  assert.ok(paused > 0);
  assert.equal(rollingTokensPerSec(samples, 80, 400 + 1000 + 1, 1000), 0);
});

test("usageTokensPerSec is output tokens over the segment duration", () => {
  assert.equal(usageTokensPerSec(278, 6200), (278 / 6200) * 1000);
  assert.equal(usageTokensPerSec(0, 1000), 0);
  assert.equal(usageTokensPerSec(10, 0), 0);
});

test("charsToTokensPerSec uses the 2-chars-per-token heuristic", () => {
  assert.equal(charsToTokensPerSec(100, 1000), 50);
});

test("shouldFreezeGeneration ends the segment on the next call, tools, idle usage, or run stop", () => {
  const base = {
    endedAt: null,
    isLoading: true,
    isLatestGenerating: true,
    hasToolCalls: false,
    grewThisTick: false,
    lastGrewAt: 1000,
    now: 1100,
    outputTokens: 0,
  };

  assert.equal(shouldFreezeGeneration(base), false);
  assert.equal(
    shouldFreezeGeneration({ ...base, isLatestGenerating: false }),
    true,
  );
  assert.equal(shouldFreezeGeneration({ ...base, hasToolCalls: true }), true);
  assert.equal(
    shouldFreezeGeneration({
      ...base,
      hasToolCalls: true,
      grewThisTick: true,
    }),
    false,
  );
  assert.equal(
    shouldFreezeGeneration({
      ...base,
      outputTokens: 40,
      now: 1000 + GENERATION_IDLE_MS,
    }),
    true,
  );
  assert.equal(
    shouldFreezeGeneration({
      ...base,
      outputTokens: 40,
      now: 1000 + GENERATION_IDLE_MS - 1,
    }),
    false,
  );
  assert.equal(shouldFreezeGeneration({ ...base, isLoading: false }), true);
  assert.equal(shouldFreezeGeneration({ ...base, endedAt: 1000 }), false);
});

test("messageHasToolCalls reads both top-level and additional_kwargs", () => {
  assert.equal(messageHasToolCalls({}), false);
  assert.equal(
    messageHasToolCalls({
      tool_calls: [{ id: "1", name: "search" }],
    }),
    true,
  );
  assert.equal(
    messageHasToolCalls({
      additional_kwargs: { tool_calls: [{ id: "1" }] },
    }),
    true,
  );
});