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

// Minimal in-memory localStorage so the override store's window-guarded functions
// run in node exactly as they do in the browser.
class MemStorage {
  private map = new Map<string, string>();
  getItem(k: string): string | null {
    return this.map.has(k) ? (this.map.get(k) as string) : null;
  }
  setItem(k: string, v: string): void {
    this.map.set(k, String(v));
  }
  removeItem(k: string): void {
    this.map.delete(k);
  }
  clear(): void {
    this.map.clear();
  }
}
const mem = new MemStorage();
(globalThis as { window?: unknown }).window = { localStorage: mem };

import {
  __resetSessionRpm,
  __setClock,
  acquire,
  defaultRpmFor,
  effectiveRpm,
  getRpmOverride,
  noteRateLimited,
  parseRetryAfterMs,
  setRpmOverride,
  type Clock,
} from "./rateLimit";

const RPM_KEY = "ytdlp-tb:ai:rpm";

function resetAll() {
  mem.clear();
  __resetSessionRpm();
}

// ─── defaultRpmFor: model-pattern matching ───

test("defaultRpmFor matches models by pattern per provider", () => {
  assert.equal(defaultRpmFor("gemini", "gemini-2.5-pro"), 5);
  assert.equal(defaultRpmFor("gemini", "gemini-1.5-pro-latest"), 5);
  assert.equal(defaultRpmFor("gemini", "gemini-2.5-flash"), 10);
  // flash-lite must win over the plain "flash" rule (its id contains "flash").
  assert.equal(defaultRpmFor("gemini", "gemini-2.5-flash-lite"), 15);
  assert.equal(defaultRpmFor("gemini", "something-else"), 10); // fallback
  assert.equal(defaultRpmFor("anthropic", "claude-haiku-4-5"), 50);
  assert.equal(defaultRpmFor("anthropic", "claude-opus-4-8"), 50);
  assert.equal(defaultRpmFor("openai", "gpt-4o-mini"), 60);
});

// ─── parseRetryAfterMs: header seconds, HTTP-date, Gemini RetryInfo ───

test("parseRetryAfterMs reads Retry-After header in seconds", () => {
  const h = new Headers({ "retry-after": "2" });
  assert.equal(parseRetryAfterMs(h), 2000);
});

test("parseRetryAfterMs reads Retry-After as an HTTP-date (against the clock)", () => {
  const base = Date.parse("2030-01-01T00:00:00Z");
  const restore = __setClock({ now: () => base, sleep: () => Promise.resolve() });
  try {
    const h = new Headers({ "retry-after": new Date(base + 5000).toUTCString() });
    assert.equal(parseRetryAfterMs(h), 5000);
  } finally {
    restore();
  }
});

test("parseRetryAfterMs reads Gemini's RetryInfo body form", () => {
  const body = JSON.stringify({
    error: {
      code: 429,
      details: [
        { "@type": "type.googleapis.com/google.rpc.RetryInfo", retryDelay: "38s" },
      ],
    },
  });
  assert.equal(parseRetryAfterMs(null, body), 38000);
});

test("parseRetryAfterMs falls back to a loose scan, else null", () => {
  assert.equal(parseRetryAfterMs(null, "please retry in 7s and try again"), 7000);
  assert.equal(parseRetryAfterMs(new Headers(), "no hint here"), null);
  assert.equal(parseRetryAfterMs(null), null);
});

// ─── limiter spacing + abort, via the injected clock ───

test("acquire spaces same-model calls by 60000/rpm (injected clock, no real time)", async () => {
  resetAll();
  let t = 1000;
  const sleeps: number[] = [];
  const clock: Clock = {
    now: () => t,
    sleep: (ms) => {
      sleeps.push(ms);
      t += ms; // advance virtual time as if we slept
      return Promise.resolve();
    },
  };
  const restore = __setClock(clock);
  try {
    // anthropic default 50 RPM → 1200ms spacing.
    await acquire("anthropic", "spacing-model");
    await acquire("anthropic", "spacing-model");
    await acquire("anthropic", "spacing-model");
    assert.deepEqual(sleeps, [1200, 1200]); // first call is free, next two paced
  } finally {
    restore();
  }
});

test("acquire aborts a pacing wait promptly", async () => {
  resetAll();
  const clock: Clock = {
    now: () => 5000, // constant, so the 2nd acquire has a nonzero wait
    sleep: (_ms, signal) =>
      new Promise<void>((_resolve, reject) => {
        if (signal?.aborted) {
          reject(new DOMException("Aborted", "AbortError"));
          return;
        }
        signal?.addEventListener(
          "abort",
          () => reject(new DOMException("Aborted", "AbortError")),
          { once: true },
        );
      }),
  };
  const restore = __setClock(clock);
  try {
    await acquire("anthropic", "abort-model"); // sets last = 5000
    const ac = new AbortController();
    const pending = acquire("anthropic", "abort-model", ac.signal);
    ac.abort();
    await assert.rejects(pending, (e: Error) => e.name === "AbortError");
  } finally {
    restore();
  }
});

// ─── step-down arithmetic ───

test("noteRateLimited steps the session RPM down by 0.6, floored, min 1", () => {
  resetAll();
  assert.equal(noteRateLimited("anthropic", "claude-x"), 30); // floor(50 * 0.6)
  assert.equal(noteRateLimited("anthropic", "claude-x"), 18); // floor(30 * 0.6)
  assert.equal(noteRateLimited("anthropic", "claude-x"), 10); // floor(18 * 0.6)
  // gemini pro starts at 5 → 3 → 1 → 1 (never below 1).
  assert.equal(noteRateLimited("gemini", "gemini-2.5-pro"), 3);
  assert.equal(noteRateLimited("gemini", "gemini-2.5-pro"), 1);
  assert.equal(noteRateLimited("gemini", "gemini-2.5-pro"), 1);
});

test("effectiveRpm precedence: override ?? session step-down ?? default", () => {
  resetAll();
  // default
  assert.equal(effectiveRpm("gemini", "gemini-2.5-pro"), 5);
  // session step-down beats default
  noteRateLimited("gemini", "gemini-2.5-pro"); // → 3
  assert.equal(effectiveRpm("gemini", "gemini-2.5-pro"), 3);
  // an explicit override beats the session step-down
  setRpmOverride("gemini", "gemini-2.5-pro", 42);
  assert.equal(effectiveRpm("gemini", "gemini-2.5-pro"), 42);
});

// ─── override store: parse-or-reset + get/set + wildcard ───

test("override store round-trips and clears", () => {
  resetAll();
  assert.equal(getRpmOverride("openai", "gpt-4o"), null);
  setRpmOverride("openai", "gpt-4o", 120);
  assert.equal(getRpmOverride("openai", "gpt-4o"), 120);
  setRpmOverride("openai", "gpt-4o", null);
  assert.equal(getRpmOverride("openai", "gpt-4o"), null);
});

test("override store resets a malformed payload rather than trapping", () => {
  resetAll();
  mem.setItem(RPM_KEY, "}{ not json");
  assert.equal(getRpmOverride("openai", "gpt-4o"), null);
  // The malformed key was cleared.
  assert.equal(mem.getItem(RPM_KEY), null);
});

test("a '*' wildcard override applies to every model", () => {
  resetAll();
  mem.setItem(RPM_KEY, JSON.stringify({ "*": 9999 }));
  assert.equal(getRpmOverride("gemini", "gemini-2.5-pro"), 9999);
  assert.equal(getRpmOverride("anthropic", "claude-haiku-4-5"), 9999);
  // An exact key still wins over the wildcard.
  setRpmOverride("anthropic", "claude-haiku-4-5", 5);
  assert.equal(getRpmOverride("anthropic", "claude-haiku-4-5"), 5);
  assert.equal(getRpmOverride("gemini", "gemini-2.5-pro"), 9999);
});