import { expect, test, vi } from "vitest";
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { CatalogDatabaseClient, CatalogPostgresAccess } from "../src/catalog/postgres-access.js";
import { ConcreteSensitivityValueSource } from "../src/catalog/sensitivity-value-source.js";
import type { CatalogColumn, CatalogTable, WorkspaceDatabase } from "../src/catalog/types.js";
import type { WorkspaceSecretStore } from "../src/workspaces/secret-store.js";
import { CATALOG_SECRET_IDS } from "../src/catalog/secrets.js";

const database = {
  id: "11111111-1111-4111-8111-111111111111",
  workspaceId: "psd-clinical",
  engine: "postgres",
  databaseName: "warehouse",
  schema: 'clinical"data',
  version: 1,
  createdAt: "2026-09-02T08:00:00Z",
  updatedAt: "2026-09-02T08:00:00Z",
  connectionStatus: "reachable",
  binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
} satisfies WorkspaceDatabase;

const table = {
  id: "22222222-2222-4222-8222-222222222222",
  databaseId: database.id,
  name: 'patient"facts',
  sourceComment: null,
  description: null,
  generatedDescription: null,
  lastSyncedDatabaseVersion: 1,
  lastSyncedAt: "2026-09-02T08:00:00Z",
  version: 1,
  createdAt: "2026-09-02T08:00:00Z",
  updatedAt: "2026-09-02T08:00:00Z",
} satisfies CatalogTable;

function column(id: string, name: string): CatalogColumn {
  return {
    id,
    tableId: table.id,
    name,
    ordinalPosition: 1,
    dataType: "text",
    isNullable: true,
    defaultExpression: null,
    primaryKeyPosition: null,
    isPrimaryKey: false,
    isForeignKey: false,
    foreignKeyCount: 0,
    sourceComment: null,
    description: null,
    generatedDescription: null,
    sensitive: false,
    lastSyncedDatabaseVersion: 1,
    lastSyncedAt: "2026-09-02T08:00:00Z",
    version: 1,
    createdAt: "2026-09-02T08:00:00Z",
    updatedAt: "2026-09-02T08:00:00Z",
  };
}

test("switches from a bounded full scan to a read-only PostgreSQL sample", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const contact = column("44444444-4444-4444-8444-444444444444", 'contact"value');
  const fullRows = Array.from({ length: 200 }, () => ({
    __value_0: "ordinary",
    __length_0: "8",
    __value_1: null,
    __length_1: null,
  }));
  const query = vi.fn(async (sql: string) => {
    if (sql.includes("TABLESAMPLE")) {
      return { rows: [{ __value_0: "sample", __length_0: 6, __value_1: "x", __length_1: 1 }] };
    }
    if (sql.startsWith("FETCH FORWARD")) return { rows: fullRows };
    return { rows: [] };
  });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  let clockCalls = 0;
  const values = new ConcreteSensitivityValueSource(access, undefined, {
    now: () => clockCalls++ < 3 ? 1_000 : 6_100,
  });
  const consumed: unknown[] = [];

  const coverage = await values.scanTable({
    database,
    table,
    columns: [note, contact],
    fullScanBudgetMs: 5_000,
    deadline: 61_000,
  }, (batch) => consumed.push(...batch), new AbortController().signal);

  expect(coverage).toEqual({ kind: "sampled", observedRows: 201 });
  expect(consumed).toContainEqual({ columnId: note.id, value: "ordinary", characterLength: 8 });
  expect(consumed).toContainEqual({ columnId: contact.id, value: null, characterLength: null });
  expect(consumed).toContainEqual({ columnId: contact.id, value: "x", characterLength: 1 });
  expect(query.mock.calls[0]).toEqual(["BEGIN TRANSACTION READ ONLY", []]);
  expect(query.mock.calls.some(([sql]) => (
    String(sql).startsWith("DECLARE sensitivity_full_scan_cursor NO SCROLL CURSOR FOR SELECT")
  ))).toBe(true);
  expect(query.mock.calls.some(([sql]) => String(sql) === (
    "FETCH FORWARD 200 FROM sensitivity_full_scan_cursor"
  ))).toBe(true);
  expect(query.mock.calls.some(([sql]) => String(sql).includes(" OFFSET "))).toBe(false);
  expect(query.mock.calls.some(([sql]) => (
    String(sql).includes('FROM "clinical""data"."patient""facts" TABLESAMPLE SYSTEM')
  ))).toBe(true);
  expect(query.mock.calls.at(-1)).toEqual(["ROLLBACK", []]);
  expect(end).toHaveBeenCalledOnce();
});

test("reports complete coverage when the final full-scan page is short", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const query = vi.fn(async (sql: string) => sql.startsWith("FETCH FORWARD")
    ? { rows: [{ __value_0: "ordinary", __length_0: 8 }] }
    : { rows: [] });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const values = new ConcreteSensitivityValueSource(access);
  const consume = vi.fn();

  const coverage = await values.scanTable({
    database,
    table,
    columns: [note],
    fullScanBudgetMs: 5_000,
    deadline: Date.now() + 60_000,
  }, consume, new AbortController().signal);

  expect(coverage).toEqual({ kind: "complete", observedRows: 1 });
  expect(query.mock.calls.filter(([sql]) => (
    String(sql) === "FETCH FORWARD 200 FROM sensitivity_full_scan_cursor"
  ))).toHaveLength(1);
  expect(consume).toHaveBeenCalledWith([
    { columnId: note.id, value: "ordinary", characterLength: 8 },
  ]);
});

test("falls back to sampling when PostgreSQL cancels the bounded full scan", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  let fullScanAttempts = 0;
  const query = vi.fn(async (sql: string) => {
    if (sql.includes("TABLESAMPLE")) {
      return { rows: [{ __value_0: "sample", __length_0: 6 }] };
    }
    if (sql.startsWith("FETCH FORWARD")) {
      fullScanAttempts += 1;
      throw Object.assign(new Error("statement timeout"), { code: "57014" });
    }
    return { rows: [] };
  });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const values = new ConcreteSensitivityValueSource(access);
  const consume = vi.fn();

  const coverage = await values.scanTable({
    database,
    table,
    columns: [note],
    fullScanBudgetMs: 5_000,
    deadline: Date.now() + 60_000,
  }, consume, new AbortController().signal);

  expect(fullScanAttempts).toBe(1);
  expect(coverage).toEqual({ kind: "sampled", observedRows: 1 });
  expect(query.mock.calls.map(([sql]) => String(sql))).toEqual(expect.arrayContaining([
    "SAVEPOINT sensitivity_full_scan",
    "ROLLBACK TO SAVEPOINT sensitivity_full_scan",
  ]));
  expect(consume).toHaveBeenCalledWith([
    { columnId: note.id, value: "sample", characterLength: 6 },
  ]);
});

test("scans a REST run_query binding without using PostgreSQL-wire access", async () => {
  const root = mkdtempSync(join(tmpdir(), "tht-sensitivity-rest-"));
  const credentialFile = join(root, "api-key");
  writeFileSync(credentialFile, "test-api-key\n", { mode: 0o600 });
  const release = vi.fn();
  const secretStore = {
    materialize: vi.fn(() => ({
      files: new Map([[CATALOG_SECRET_IDS.apiKey, credentialFile]]),
      release,
    })),
  } as unknown as WorkspaceSecretStore;
  const fetchMock = vi.fn(async () => new Response(JSON.stringify([
    { __value_0: "mario.rossi@example.it", __length_0: 23 },
  ]), { status: 200, headers: { "content-type": "application/json" } }));
  vi.stubGlobal("fetch", fetchMock);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => { throw new Error("PostgreSQL access must not be used"); }),
  };
  const values = new ConcreteSensitivityValueSource(access, secretStore);
  const restDatabase: WorkspaceDatabase = {
    ...database,
    binding: {
      transport: "rest_api",
      baseUrl: "https://dwh.example.test/root/",
      restPath: "/health",
      restAuth: "x-api-key",
    },
  };
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const consume = vi.fn();

  try {
    await expect(values.scanTable({
      database: restDatabase,
      table,
      columns: [note],
      fullScanBudgetMs: 5_000,
      deadline: Date.now() + 60_000,
    }, consume, new AbortController().signal)).resolves.toEqual({
      kind: "complete",
      observedRows: 1,
    });
    expect(access.connect).not.toHaveBeenCalled();
    expect(fetchMock).toHaveBeenCalledWith(
      "https://dwh.example.test/root/rpc/run_query",
      expect.objectContaining({
        method: "POST",
        headers: { "content-type": "application/json", "x-api-key": "test-api-key" },
      }),
    );
    const body = JSON.parse(String(fetchMock.mock.calls[0]![1]!.body));
    expect(body.query_text).toContain('FROM "clinical""data"."patient""facts" LIMIT 200 OFFSET 0');
    expect(consume).toHaveBeenCalledWith([
      { columnId: note.id, value: "mario.rossi@example.it", characterLength: 23 },
    ]);
    expect(release).toHaveBeenCalledOnce();
  } finally {
    vi.unstubAllGlobals();
    rmSync(root, { recursive: true, force: true });
  }
});

test("keeps multi-request REST scans conservative without a source transaction", async () => {
  const root = mkdtempSync(join(tmpdir(), "tht-sensitivity-rest-pages-"));
  const credentialFile = join(root, "api-key");
  writeFileSync(credentialFile, "test-api-key\n", { mode: 0o600 });
  const secretStore = {
    materialize: vi.fn(() => ({
      files: new Map([[CATALOG_SECRET_IDS.apiKey, credentialFile]]),
      release: vi.fn(),
    })),
  } as unknown as WorkspaceSecretStore;
  const fetchMock = vi.fn()
    .mockResolvedValueOnce(new Response(JSON.stringify([
      { __value_0: "ordinary", __length_0: 8 },
    ]), { status: 200 }))
    .mockResolvedValueOnce(new Response(JSON.stringify([]), { status: 200 }));
  vi.stubGlobal("fetch", fetchMock);
  const values = new ConcreteSensitivityValueSource({
    connect: vi.fn(async () => { throw new Error("PostgreSQL access must not be used"); }),
  }, secretStore, { batchRows: 1 });
  const restDatabase: WorkspaceDatabase = {
    ...database,
    binding: {
      transport: "rest_api",
      baseUrl: "https://dwh.example.test/root",
      restPath: "/health",
      restAuth: "x-api-key",
    },
  };

  try {
    await expect(values.scanTable({
      database: restDatabase,
      table,
      columns: [column("33333333-3333-4333-8333-333333333333", "note")],
      fullScanBudgetMs: 5_000,
      deadline: Date.now() + 60_000,
    }, vi.fn(), new AbortController().signal)).resolves.toEqual({
      kind: "sampled",
      observedRows: 1,
    });
    expect(fetchMock).toHaveBeenCalledTimes(2);
  } finally {
    vi.unstubAllGlobals();
    rmSync(root, { recursive: true, force: true });
  }
});

test("does not start a PostgreSQL transaction when connecting consumed the run deadline", async () => {
  const query = vi.fn(async () => ({ rows: [] }));
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const now = vi.fn()
    .mockReturnValueOnce(1_000)
    .mockReturnValue(61_000);
  const values = new ConcreteSensitivityValueSource(access, undefined, { now });

  await expect(values.scanTable({
    database,
    table,
    columns: [column("33333333-3333-4333-8333-333333333333", "note")],
    fullScanBudgetMs: 5_000,
    deadline: 60_000,
  }, vi.fn(), new AbortController().signal)).resolves.toEqual({
    kind: "sampled",
    observedRows: 0,
  });
  expect(query).not.toHaveBeenCalled();
  expect(end).toHaveBeenCalledOnce();
});