import { expect, test, vi } from "vitest";
import {
  SensitivityAnalysisInterruptedError,
  SensitivityAnalysisService,
} from "../src/catalog/sensitivity-analysis-service.js";
import { SensitivityAnalysisRunner } from "../src/catalog/sensitivity-analysis-runner.js";
import type { SensitivityClassifier } from "../src/catalog/sensitivity-classifier.js";
import type {
  CatalogRepository,
  SensitivityAnalysisRun,
  WorkspaceDatabase,
} from "../src/catalog/types.js";

const database = {
  id: "11111111-1111-4111-8111-111111111111",
  workspaceId: "psd-clinical",
  engine: "postgres",
  databaseName: "warehouse",
  schema: "public",
  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 running: SensitivityAnalysisRun = {
  id: "22222222-2222-4222-8222-222222222222",
  databaseId: database.id,
  scope: "all",
  engine: "local",
  modelId: null,
  policyVersion: "sensitivity-v1",
  status: "running",
  total: 0,
  suggestedSensitive: 0,
  suggestedNonSensitive: 0,
  unknown: 0,
  inputTokens: 0,
  cacheReadTokens: 0,
  outputTokens: 0,
  createdAt: "2026-09-02T08:00:00Z",
  startedAt: "2026-09-02T08:00:00Z",
  updatedAt: "2026-09-02T08:00:00Z",
  finishedAt: null,
  errorSummary: null,
};

test("stops catalog selection when the request expires during a catalog read", async () => {
  const controller = new AbortController();
  const listTables = vi.fn();
  const repository = {
    get: vi.fn(async () => {
      controller.abort();
      return database;
    }),
    listTables,
  } as unknown as CatalogRepository;
  const classifier = { assessTable: vi.fn() } as unknown as SensitivityClassifier;
  const analysis = new SensitivityAnalysisService(repository, classifier);

  await expect(analysis.analyze(
    database.id,
    "all",
    [],
    controller.signal,
  )).rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
  expect(listTables).not.toHaveBeenCalled();
  expect(classifier.assessTable).not.toHaveBeenCalled();
});

test("marks a created run interrupted if the request deadline expires during persistence", async () => {
  const controller = new AbortController();
  const update = vi.fn(async (_runId: string, changes: Partial<SensitivityAnalysisRun>) => ({
    ...running,
    ...changes,
  }));
  const repository = {
    get: vi.fn(async () => database),
    createSensitivityAnalysisRun: vi.fn(async () => {
      controller.abort();
      return running;
    }),
    updateSensitivityAnalysisRun: update,
    appendSensitivityAnalysisEvent: vi.fn(async () => undefined),
  } as unknown as CatalogRepository;
  const analysis = { analyze: vi.fn() } as unknown as SensitivityAnalysisService;
  const runner = new SensitivityAnalysisRunner(repository, analysis);

  await expect(runner.run(database.id, "all", [], controller.signal))
    .rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
  expect(analysis.analyze).not.toHaveBeenCalled();
  expect(update).toHaveBeenCalledWith(running.id, expect.objectContaining({
    status: "interrupted",
    total: 0,
    unknown: 0,
    errorSummary: "Local sensitivity analysis reached its time limit.",
  }));
});