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

function request(columns: readonly CatalogColumn[], overrides: Record<string, unknown> = {}) {
  return {
    database,
    table,
    columns,
    valuesPerColumn: 300,
    sampleOffset: 0,
    sampleSeed: 37,
    queryTimeoutMs: 5_000,
    fullScanThreshold: 1_000,
    ...overrides,
  };
}

test("uses bounded read-only PostgreSQL sampling for tables above 1,000 rows", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const contact = column("44444444-4444-4444-8444-444444444444", 'contact"value');
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("SELECT 1 AS __present")) {
      return { rows: Array.from({ length: 1_001 }, () => ({ __present: 1 })) };
    }
    if (sql.startsWith("WITH sampled")) {
      return { rows: [
        { __column_index: 0, __value: "ordinary", __length: "8" },
        { __column_index: 1, __value: "mario.rossi@example.it", __length: 23 },
      ] };
    }
    return { rows: [] };
  });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const consume = vi.fn();

  await expect(new ConcreteSensitivityValueSource(access).scanTable(
    request([note, contact]),
    consume,
    new AbortController().signal,
  )).resolves.toEqual({ kind: "sampled", observedValues: 2 });

  expect(query.mock.calls[0]).toEqual(["BEGIN TRANSACTION READ ONLY", []]);
  expect(query).toHaveBeenCalledWith("SELECT set_config('statement_timeout', $1, true)", ["5000ms"]);
  const sampleSql = query.mock.calls.map(([sql]) => String(sql)).find((sql) => sql.startsWith("WITH sampled"));
  expect(sampleSql).toContain('FROM "clinical""data"."patient""facts" TABLESAMPLE SYSTEM (30)');
  expect(sampleSql).toContain("REPEATABLE (37)");
  expect(sampleSql).toContain("LIMIT 3000 OFFSET 0");
  expect(sampleSql).toContain("CROSS JOIN LATERAL");
  expect(sampleSql).toContain("WHERE __rank <= 300");
  expect(consume).toHaveBeenCalledWith([
    { columnId: note.id, value: "ordinary", characterLength: 8 },
    { columnId: contact.id, value: "mario.rossi@example.it", characterLength: 23 },
  ]);
  expect(query.mock.calls.at(-1)).toEqual(["ROLLBACK", []]);
  expect(end).toHaveBeenCalledOnce();
});

test("fully scans a table when the 1,001-row probe proves it is small", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("SELECT 1 AS __present")) return { rows: [{ __present: 1 }] };
    if (sql.startsWith("WITH sampled")) {
      return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
    }
    return { rows: [] };
  });
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
  };
  const consume = vi.fn();

  await expect(new ConcreteSensitivityValueSource(access).scanTable(
    request([note]),
    consume,
    new AbortController().signal,
  )).resolves.toEqual({ kind: "complete", observedValues: 1 });

  const valueSql = query.mock.calls.map(([sql]) => String(sql)).find((sql) => sql.startsWith("WITH sampled"));
  expect(valueSql).not.toContain("TABLESAMPLE");
  expect(valueSql).toContain("WHERE __rank <= 1000");
  expect(consume).toHaveBeenCalledWith([
    { columnId: note.id, value: "ordinary", characterLength: 8 },
  ]);
});

test("falls back to sampling when the small-table probe reaches its query timeout", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("SELECT 1 AS __present")) {
      throw Object.assign(new Error("statement timeout"), { code: "57014" });
    }
    if (sql.startsWith("WITH sampled")) {
      return { rows: [{ __column_index: 0, __value: "sample", __length: 6 }] };
    }
    return { rows: [] };
  });
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
  };
  const consume = vi.fn();

  await expect(new ConcreteSensitivityValueSource(access).scanTable(
    request([note]),
    consume,
    new AbortController().signal,
  )).resolves.toEqual({ kind: "sampled", observedValues: 1 });
  expect(query.mock.calls.map(([sql]) => String(sql))).toContain(
    "ROLLBACK TO SAVEPOINT sensitivity_scan_1",
  );
});

test("limits each source query to at most 25 columns", async () => {
  const columns = Array.from({ length: 26 }, (_, index) => column(
    `00000000-0000-4000-8000-${(index + 1).toString().padStart(12, "0")}`,
    `attribute_${index + 1}`,
  ));
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("SELECT 1 AS __present")) {
      return { rows: Array.from({ length: 1_001 }, () => ({ __present: 1 })) };
    }
    if (sql.startsWith("WITH sampled")) {
      return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
    }
    return { rows: [] };
  });
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
  };

  await new ConcreteSensitivityValueSource(access).scanTable(
    request(columns),
    vi.fn(),
    new AbortController().signal,
  );

  expect(query.mock.calls.filter(([sql]) => String(sql).startsWith("WITH sampled"))).toHaveLength(2);
});

test("scans a REST run_query binding without 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([
    { __column_index: 0, __value: "mario.rossi@example.it", __length: 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(request([note], {
      database: restDatabase,
      fullScanThreshold: undefined,
    }), consume, new AbortController().signal)).resolves.toEqual({
      kind: "sampled",
      observedValues: 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" TABLESAMPLE SYSTEM (30)');
    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("falls back to a sequential bounded sample when randomized sampling times out", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("WITH sampled") && sql.includes("TABLESAMPLE")) {
      throw Object.assign(new Error("raw source detail"), { code: "57014" });
    }
    if (sql.startsWith("WITH sampled")) {
      return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
    }
    return { rows: [] };
  });
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
  };

  await expect(new ConcreteSensitivityValueSource(access).scanTable(
    request([note], { fullScanThreshold: undefined }),
    vi.fn(),
    new AbortController().signal,
  )).resolves.toEqual({ kind: "sampled", observedValues: 1 });
  expect(query.mock.calls.filter(([sql]) => String(sql).startsWith("WITH sampled"))).toHaveLength(2);
});

test("fails explicitly when both randomized and sequential sample queries time out", async () => {
  const note = column("33333333-3333-4333-8333-333333333333", "note");
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("WITH sampled")) {
      throw Object.assign(new Error("raw source detail"), { code: "57014" });
    }
    return { rows: [] };
  });
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
  };

  await expect(new ConcreteSensitivityValueSource(access).scanTable(
    request([note], { fullScanThreshold: undefined }),
    vi.fn(),
    new AbortController().signal,
  )).rejects.toEqual(new CatalogConnectorError("Sensitivity sample query timed out"));
});