import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
import { readFileSync } from "node:fs";
import { fakeD1 } from "../metering/fake-d1.test-helper";

let d1 = fakeD1();
const kvStore = new Map<string, string>();
const fakeEnv: Record<string, unknown> = {
  get DB() {
    return d1.binding;
  },
  SECRETS_KEY: "unit-test-secrets-key-0123456789abcdef",
  PUBLIC_URL: "https://forge.test",
  INTEGRATIONS_ALLOW_LOOPBACK: "false",
  MERGEFIX_GATEWAY_ID: "default",
  KV: {
    get: async (key: string) => kvStore.get(key) ?? null,
    put: async (key: string, value: string) => void kvStore.set(key, value),
  },
};
mock.module("cloudflare:workers", () => ({ env: fakeEnv, waitUntil: () => undefined }));
const { sharedScopes } = await import("../test-shared-scopes");
mock.module("../scoped-db", () => ({
  accountDb: (id: string) => sharedScopes(d1.db).accountDb(id),
  repoDb: (id: string) => sharedScopes(d1.db).repoDb(id),
  stmt: (sql: string, ...params: (string | number | null)[]) => ({ sql, params }),
}));

const {
  validateEndpoint,
  saveEndpoint,
  endpointSummary,
  resolveAiEndpoint,
  resolveAiEndpointSummary,
  deleteEndpoint,
  resealLegacyEndpoints,
} = await import("./endpoint");
const { chatCompletion, endpointUrl } = await import("../mergefix/gateway");

const migration = readFileSync(new URL("../../../migrations/0032_ai_endpoints.sql", import.meta.url), "utf8");
const capMigration = readFileSync(new URL("../../../migrations/0037_byok_token_cap.sql", import.meta.url), "utf8");
const realFetch = globalThis.fetch;

beforeEach(() => {
  d1 = fakeD1();
  d1.db.exec(`CREATE TABLE account_ai_keys (account_id TEXT, provider TEXT, model TEXT, key_ciphertext TEXT,
    key_iv TEXT, updated_by TEXT, updated_at INTEGER);`);
  d1.db.exec(migration);
  d1.db.exec(capMigration);
  d1.db.exec(`
    INSERT INTO accounts (id, name, kind) VALUES ('u1', 'alice', 'user'), ('u2', 'bob', 'user');
    INSERT INTO repos VALUES ('r1', 'u1', 'web');
  `);
  fakeEnv.INTEGRATIONS_ALLOW_LOOPBACK = "false";
});

afterEach(() => {
  globalThis.fetch = realFetch;
});

const provider = { kind: "provider", provider: "openai", model: "gpt-4.1-mini", apiKey: "sk-account-key" } as const;
const custom = {
  kind: "openai",
  model: "llama3",
  baseUrl: "https://llm.example.com/v1/",
  apiKey: "sk-repo-key-1",
} as const;

describe("SSRF validation", () => {
  test.each([
    "http://llm.example.com/v1",
    "https://10.0.0.1/v1",
    "https://127.0.0.1/v1",
    "https://localhost/v1",
    "https://[::1]/v1",
    "https://169.254.169.254/latest",
    "https://intranet/v1",
    "https://user:pass@llm.example.com/v1",
    "https://forge.test/api",
    "file:///etc/passwd",
  ])("rejects %s", (baseUrl) => {
    expect(() => validateEndpoint({ ...custom, baseUrl })).toThrow();
  });

  test("accepts public https and normalizes trailing slashes", () => {
    expect(validateEndpoint(custom).baseUrl).toBe("https://llm.example.com/v1");
  });

  test("loopback is allowed only with INTEGRATIONS_ALLOW_LOOPBACK on a loopback forge", () => {
    const local = { ...custom, baseUrl: "http://localhost:11434/v1" };
    expect(() => validateEndpoint(local)).toThrow();
    fakeEnv.INTEGRATIONS_ALLOW_LOOPBACK = "true";
    expect(() => validateEndpoint(local)).toThrow();
    fakeEnv.PUBLIC_URL = "http://localhost:3000";
    fakeEnv.GITFOR_LOCAL = "true";
    try {
      expect(validateEndpoint(local).baseUrl).toBe("http://localhost:11434/v1");
    } finally {
      fakeEnv.PUBLIC_URL = "https://forge.test";
      fakeEnv.GITFOR_LOCAL = undefined;
      fakeEnv.INTEGRATIONS_ALLOW_LOOPBACK = "false";
    }
  });

  test("Cloudflare gateway ids are validated and the URL is fixed", async () => {
    expect(() =>
      validateEndpoint({ ...provider, kind: "cf-gateway", cfAccountId: "../../x", cfGatewayId: "gw" }),
    ).toThrow();
    const valid = validateEndpoint({
      ...provider,
      kind: "cf-gateway",
      cfAccountId: "0123456789abcdef0123456789ABCDEF",
      cfGatewayId: "my-gw",
    });
    expect(await endpointUrl({ ...valid, scope: "account", baseUrl: null } as Parameters<typeof endpointUrl>[0])).toBe(
      "https://gateway.ai.cloudflare.com/v1/0123456789abcdef0123456789abcdef/my-gw/compat/chat/completions",
    );
  });

  test("a hostname that resolves to a private address is refused at call time", async () => {
    let called = false;
    globalThis.fetch = (async (input: RequestInfo | URL) => {
      const url = String(input instanceof Request ? input.url : input);
      if (url.startsWith("https://cloudflare-dns.com/")) {
        const type = new URL(url).searchParams.get("type");
        return Response.json({ Status: 0, Answer: type === "A" ? [{ type: 1, data: "10.0.0.5" }] : [] });
      }
      called = true;
      return Response.json({});
    }) as typeof fetch;
    await saveEndpoint("account", "u1", custom, "u1");
    const endpoint = await resolveAiEndpoint({ accountId: "u1" });
    await expect(
      chatCompletion({ model: "@cf/x", byok: endpoint, metadata: { account_id: "u1", feature: "t" } }, []),
    ).rejects.toThrow(/non-public/);
    expect(called).toBe(false);
  });
});

describe("resolver", () => {
  test("repo overrides account, and keys are never in summaries", async () => {
    await saveEndpoint("account", "u1", provider, "u1");
    expect((await resolveAiEndpoint({ accountId: "u1", repoId: "r1" }))?.apiKey).toBe("sk-account-key");
    const saved = await saveEndpoint("repo", "r1", custom, "u1");
    expect(JSON.stringify(saved)).not.toContain("sk-repo-key-1");
    expect(JSON.stringify(await endpointSummary("repo", "r1"))).not.toContain("sk-repo-key-1");
    const resolved = await resolveAiEndpoint({ accountId: "u1", repoId: "r1" });
    expect(resolved?.scope).toBe("repo");
    expect(resolved?.apiKey).toBe("sk-repo-key-1");
    const row = d1.db.query("SELECT key_ciphertext FROM ai_endpoints WHERE scope = 'repo'").get() as {
      key_ciphertext: string;
    };
    expect(row.key_ciphertext).not.toContain("sk-repo-key-1");
    await deleteEndpoint("repo", "r1");
    expect((await resolveAiEndpoint({ accountId: "u1", repoId: "r1" }))?.scope).toBe("account");
  });

  test("a repo's endpoint never pays for another account's calls", async () => {
    await saveEndpoint("repo", "r1", custom, "u1");
    expect(await resolveAiEndpoint({ accountId: "u2", repoId: "r1" })).toBeNull();
    expect(await resolveAiEndpointSummary({ accountId: "u2", repoId: "r1" })).toBeNull();
  });

  test("a sealed key can't be swapped onto another owner", async () => {
    await saveEndpoint("account", "u1", provider, "u1");
    d1.db.exec(`INSERT INTO ai_endpoints (scope, owner_id, kind, provider, model, key_ciphertext, key_iv, updated_at)
      SELECT 'account', 'u2', kind, provider, model, key_ciphertext, key_iv, updated_at FROM ai_endpoints WHERE owner_id = 'u1'`);
    await expect(resolveAiEndpoint({ accountId: "u2" })).rejects.toThrow();
  });

  test("keys migrated from account_ai_keys are only usable once the cron re-seals them to their owner", async () => {
    const { encryptSecret } = await import("../crypto");
    const legacy = await encryptSecret(String(fakeEnv.SECRETS_KEY), "sk-legacy-key");
    d1.db
      .query(
        `INSERT INTO ai_endpoints (scope, owner_id, kind, provider, model, key_ciphertext, key_iv, scoped_seal, updated_at)
         VALUES ('account', 'u1', 'provider', 'openai', 'gpt', ?, ?, 0, 0)`,
      )
      .run(legacy.ciphertext, legacy.iv);
    await expect(resolveAiEndpoint({ accountId: "u1" })).rejects.toThrow("upgraded");
    await resealLegacyEndpoints();
    const row = d1.db.query("SELECT key_ciphertext, scoped_seal FROM ai_endpoints WHERE owner_id = 'u1'").get() as {
      key_ciphertext: string;
      scoped_seal: number;
    };
    expect(row.scoped_seal).toBe(1);
    expect(row.key_ciphertext).not.toBe(legacy.ciphertext);
    expect((await resolveAiEndpoint({ accountId: "u1" }))?.apiKey).toBe("sk-legacy-key");
    await resealLegacyEndpoints();
    expect((await resolveAiEndpoint({ accountId: "u1" }))?.apiKey).toBe("sk-legacy-key");
  });

  test("a key with a monthly token cap stops being used once the month's BYOK tokens reach it", async () => {
    await saveEndpoint("account", "u1", { ...provider, monthlyTokenCap: 100 }, "u1");
    expect((await resolveAiEndpoint({ accountId: "u1" }))?.monthlyTokenCap).toBe(100);
    d1.db
      .query(
        "INSERT OR REPLACE INTO usage_periods (account_id, meter, period, quantity, updated_at) VALUES ('u1', 'ai.byok_tokens', ?, 100, 0)",
      )
      .run(new Date().toISOString().slice(0, 7));
    await expect(resolveAiEndpoint({ accountId: "u1" })).rejects.toThrow("cap");
    expect(() => validateEndpoint({ ...provider, monthlyTokenCap: 0 })).toThrow("token cap");
  });

  test("deleting the repo removes its endpoint", async () => {
    await saveEndpoint("repo", "r1", custom, "u1");
    d1.db.exec("DELETE FROM repos WHERE id = 'r1'");
    expect(await endpointSummary("repo", "r1")).toBeNull();
  });
});