import { expect, test, vi } from "vitest";
import { DescriptionGenerationWorker } from "../src/catalog/description-generation-worker.js";
import type { DescriptionSourceSampler } from "../src/catalog/description-source-sampler.js";
import { MemoryCatalogRepository } from "../src/catalog/memory-repository.js";
import type { MetadataGenerationModels } from "../src/catalog/metadata-generation-models.js";
import { ModelCompletionCancelledError } from "../src/catalog/model-completer.js";
import type { ModelCompleter, ModelCompletionRequest } from "../src/catalog/model-completer.js";
import { CatalogOperationCoordinator } from "../src/catalog/operation-coordinator.js";
import type { WorkspaceRegistry } from "../src/workspaces/registry.js";

test("serializes Unlock with Start so stale recovery cannot release a new reservation", async () => {
  let lookupStarted!: () => void;
  const started = new Promise<void>((resolve) => { lookupStarted = resolve; });
  let releaseLookup!: () => void;
  const gate = new Promise<void>((resolve) => { releaseLookup = resolve; });
  const repository = {
    getActiveDescriptionGenerationRun: vi.fn(async () => {
      lookupStarted();
      await gate;
      return undefined;
    }),
  } as unknown as MemoryCatalogRepository;
  const resolveModel = vi.fn();
  const worker = new DescriptionGenerationWorker(
    repository,
    {} as WorkspaceRegistry,
    {
      catalog: () => ({ models: [], default: "" }),
      resolve: resolveModel,
    } as MetadataGenerationModels,
    {} as ModelCompleter,
    new CatalogOperationCoordinator(),
    { sample: vi.fn(async () => []) },
  );

  const unlocking = worker.unlock();
  await started;
  await expect(worker.start(
    "11111111-1111-4111-8111-111111111111",
    "openai-mini",
    "missing",
    [],
  )).rejects.toThrow("already active");
  expect(resolveModel).not.toHaveBeenCalled();

  releaseLookup();
  await expect(unlocking).resolves.toBeUndefined();
});

test("exposes an awaitable background job and absorbs provider promise rejection", async () => {
  const repository = new MemoryCatalogRepository();
  const database = await repository.create({
    workspaceId: "psd-clinical",
    engine: "postgres",
    databaseName: "warehouse",
    schema: "datawarehouse",
    binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
  });
  await repository.applySchemaSync(database.id, database.version, "all", [], {
    schemaVersion: 1,
    capabilities: { tables: "available", columns: "available", relationships: "available" },
    tables: [{ name: "patients", sourceComment: null }],
    columns: [{
      tableName: "patients",
      name: "birth_date",
      ordinalPosition: 1,
      dataType: "date",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }],
    relationships: [],
  });
  const table = (await repository.listTables(database.id))[0]!;
  const column = (await repository.listColumns(database.id, table.id))[0]!;
  let rejectCompletion!: (error: Error) => void;
  const pendingCompletion = new Promise<string>((_resolve, reject) => { rejectCompletion = reject; });
  const completer: ModelCompleter = {
    complete: vi.fn(async () => await pendingCompletion),
  };
  const models: MetadataGenerationModels = {
    catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
    resolve: () => ({
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    }),
  };
  const operations = new CatalogOperationCoordinator();
  const sourceSampler: DescriptionSourceSampler = {
    sample: vi.fn(async () => []),
  };
  const worker = new DescriptionGenerationWorker(
    repository,
    {
      read: vi.fn(async () => ({
        workspace: { workspace: { language: "it" } },
        revision: {},
      })),
    } as unknown as WorkspaceRegistry,
    models,
    completer,
    operations,
    sourceSampler,
  );

  const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
  let settled = false;
  const waiting = worker.waitForRun(run.id).then(() => { settled = true; });
  await new Promise((resolve) => setTimeout(resolve, 0));
  expect(settled).toBe(false);

  rejectCompletion(new Error("test-provider-secret private prompt raw response"));
  await expect(waiting).resolves.toBeUndefined();
  expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
    status: "completed_with_errors",
    failed: 1,
    errorSummary: "Description generation completed with errors.",
  });
  const events = await repository.listDescriptionGenerationEvents(run.id);
  expect(JSON.stringify(events)).not.toMatch(/test-provider-secret|private prompt|raw response/);

  const release = operations.reserve(database.id);
  release();

  await expect(worker.start(
    database.id,
    "openai-mini",
    "selected_columns",
    [],
  )).rejects.toThrow("at least one target ID is required");
});

test("marks an active run interrupted when the backend worker stops", async () => {
  const repository = new MemoryCatalogRepository();
  const database = await repository.create({
    workspaceId: "psd-clinical",
    engine: "postgres",
    databaseName: "warehouse",
    schema: "datawarehouse",
    binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
  });
  await repository.applySchemaSync(database.id, database.version, "all", [], {
    schemaVersion: 1,
    capabilities: { tables: "available", columns: "available", relationships: "available" },
    tables: [{ name: "patients", sourceComment: null }],
    columns: [{
      tableName: "patients",
      name: "status",
      ordinalPosition: 1,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }],
    relationships: [],
  });
  const table = (await repository.listTables(database.id))[0]!;
  const column = (await repository.listColumns(database.id, table.id))[0]!;
  const completer: ModelCompleter = {
    complete: vi.fn(async (request) => await new Promise<string>((_resolve, reject) => {
      const cancel = () => reject(new ModelCompletionCancelledError());
      if (request.signal.aborted) cancel();
      else request.signal.addEventListener("abort", cancel, { once: true });
    })),
  };
  const worker = new DescriptionGenerationWorker(
    repository,
    {
      read: vi.fn(async () => ({
        workspace: { workspace: { language: "it" } },
        revision: {},
      })),
    } as unknown as WorkspaceRegistry,
    {
      catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
      resolve: () => ({
        id: "openai-mini",
        provider: "openai",
        model: "gpt-4.1-mini",
        apiKeyEnv: "OPENAI_API_KEY",
        apiKey: "test-provider-secret",
      }),
    },
    completer,
    new CatalogOperationCoordinator(),
    { sample: vi.fn(async () => []) },
  );

  const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
  await vi.waitFor(() => expect(completer.complete).toHaveBeenCalledOnce());
  await worker.stop();

  expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
    status: "interrupted",
    errorSummary: "Description generation was interrupted by backend shutdown.",
  });
  expect(await repository.listDescriptionGenerationEvents(run.id)).toContainEqual(
    expect.objectContaining({
      level: "warning",
      message: "Description generation was interrupted by backend shutdown.",
    }),
  );
});

test("adds only bounded transient source samples to the model request", async () => {
  const repository = new MemoryCatalogRepository();
  const database = await repository.create({
    workspaceId: "psd-clinical",
    engine: "postgres",
    databaseName: "warehouse",
    schema: "datawarehouse",
    binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
  });
  await repository.applySchemaSync(database.id, database.version, "all", [], {
    schemaVersion: 1,
    capabilities: { tables: "available", columns: "available", relationships: "available" },
    tables: [{ name: "patients", sourceComment: null }],
    columns: [{
      tableName: "patients",
      name: "patient_email",
      ordinalPosition: 1,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }, {
      tableName: "patients",
      name: "ward",
      ordinalPosition: 2,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }],
    relationships: [],
  });
  const table = (await repository.listTables(database.id))[0]!;
  const columns = await repository.listColumns(database.id, table.id);
  const column = columns.find((candidate) => candidate.name === "patient_email")!;
  const ward = columns.find((candidate) => candidate.name === "ward")!;
  await repository.updateColumnMetadata(
    database.id,
    table.id,
    column.id,
    column.version,
    column.description,
    column.generatedDescription,
    true,
  );
  const sampleSecret = "real.patient@hospital.invalid";
  const sourceSampler: DescriptionSourceSampler = {
    sample: vi.fn(async () => [{
      targetId: column.id,
      tableName: table.name,
      rows: [
        { fields: [{ name: column.name, value: sampleSecret }] },
        { fields: [{ name: column.name, value: "row-2" }] },
        { fields: [{ name: column.name, value: "row-3" }] },
      ],
      representativeValues: [{
        column: column.name,
        values: [sampleSecret, sampleSecret, "two", "three"],
      }],
    }, {
      targetId: ward.id,
      tableName: table.name,
      rows: [
        { fields: [{ name: ward.name, value: "row-4" }] },
        { fields: [{ name: ward.name, value: "row-5" }] },
        { fields: [{ name: ward.name, value: "row-6" }] },
        { fields: [{ name: ward.name, value: "row-7" }] },
        { fields: [{ name: ward.name, value: "row-8" }] },
        { fields: [{ name: ward.name, value: "row-9-must-be-omitted" }] },
      ],
      representativeValues: [{
        column: ward.name,
        values: ["ward-1", "ward-2", "ward-3", "ward-4", "ward-5", "ward-6-must-be-omitted"],
      }],
    }]),
  };
  const completer: ModelCompleter = {
    complete: vi.fn(async () => JSON.stringify({
      results: [{
        targetId: column.id,
        outcome: "generated",
        description: "Stato amministrativo del paziente.",
      }, {
        targetId: ward.id,
        outcome: "generated",
        description: "Reparto associato al paziente.",
      }],
    })),
  };
  const models: MetadataGenerationModels = {
    catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
    resolve: () => ({
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    }),
  };
  const worker = new DescriptionGenerationWorker(
    repository,
    {
      read: vi.fn(async () => ({
        workspace: { workspace: { language: "it" } },
        revision: {},
      })),
    } as unknown as WorkspaceRegistry,
    models,
    completer,
    new CatalogOperationCoordinator(),
    sourceSampler,
  );

  const run = await worker.start(
    database.id,
    "openai-mini",
    "selected_columns",
    [column.id, ward.id],
  );
  await worker.waitForRun(run.id);

  expect(sourceSampler.sample).toHaveBeenCalledWith(
    expect.objectContaining({ id: database.id, binding: database.binding }),
    [
      { targetId: column.id, tableName: table.name, columnNames: [] },
      { targetId: ward.id, tableName: table.name, columnNames: [ward.name] },
    ],
    expect.any(AbortSignal),
  );
  const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
  expect(request.messages[0]?.content).toContain("untrusted");
  const userMessage = request.messages[1]!.content;
  const context = JSON.parse(userMessage.slice(userMessage.indexOf("\n") + 1));
  const sampledRows = context.targets.flatMap(
    (targetContext: { sourceSample?: { rows: unknown[] } }) => targetContext.sourceSample?.rows ?? [],
  );
  const representativeValues = context.targets.flatMap(
    (targetContext: { sourceSample?: { representativeValues: Array<{ values: unknown[] }> } }) => (
      targetContext.sourceSample?.representativeValues.flatMap((entry) => entry.values) ?? []
    ),
  );
  expect(sampledRows).toHaveLength(10);
  expect(representativeValues).toHaveLength(10);
  expect(context.targets[0].sourceSample.rows).toHaveLength(5);
  expect(context.targets[1].sourceSample.rows).toHaveLength(5);
  expect(context.targets[0].sourceSample.representativeValues).toEqual([{
    column: column.name,
    values: [
      "marta.rossi@example.com",
      "luca.bianchi@example.com",
      "elena.conti@example.com",
      "paolo.romano@example.com",
      "giulia.ferrari@example.com",
    ],
  }]);
  expect(context.targets[1].sourceSample.representativeValues).toEqual([{
    column: ward.name,
    values: ["ward-1", "ward-2", "ward-3", "ward-4", "ward-5"],
  }]);
  expect(userMessage).not.toContain(sampleSecret);
  expect(userMessage).not.toMatch(/synthetic|fake|fittizi/i);
  expect(userMessage).toContain("marta.rossi@example.com");
  expect(userMessage).not.toMatch(
    /row-9-must-be-omitted|ward-6-must-be-omitted/,
  );

  const persisted = JSON.stringify({
    run: await repository.getDescriptionGenerationRun(run.id),
    events: await repository.listDescriptionGenerationEvents(run.id),
    database: await repository.get(database.id),
    table: await repository.getTable(database.id, table.id),
    column: await repository.getColumn(database.id, table.id, column.id),
    ward: await repository.getColumn(database.id, table.id, ward.id),
  });
  expect(persisted).not.toContain(sampleSecret);
});

test("gives every sensitive column synthetic context without consuming the real sample budget", async () => {
  const repository = new MemoryCatalogRepository();
  const database = await repository.create({
    workspaceId: "psd-clinical",
    engine: "postgres",
    databaseName: "warehouse",
    schema: "datawarehouse",
    binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
  });
  await repository.applySchemaSync(database.id, database.version, "all", [], {
    schemaVersion: 1,
    capabilities: { tables: "available", columns: "available", relationships: "available" },
    tables: [{ name: "patients", sourceComment: null }],
    columns: [{
      tableName: "patients",
      name: "patient_email",
      ordinalPosition: 1,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }, {
      tableName: "patients",
      name: "patient_phone",
      ordinalPosition: 2,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }, {
      tableName: "patients",
      name: "ward",
      ordinalPosition: 3,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }],
    relationships: [],
  });
  const table = (await repository.listTables(database.id))[0]!;
  const columns = await repository.listColumns(database.id, table.id);
  const email = columns.find((column) => column.name === "patient_email")!;
  const phone = columns.find((column) => column.name === "patient_phone")!;
  const ward = columns.find((column) => column.name === "ward")!;
  for (const column of [email, phone]) {
    await repository.updateColumnMetadata(
      database.id,
      table.id,
      column.id,
      column.version,
      column.description,
      column.generatedDescription,
      true,
    );
  }
  const wardValues = ["ward-a", "ward-b", "ward-c", "ward-d", "ward-e"];
  const sourceSampler: DescriptionSourceSampler = {
    sample: vi.fn(async (_database, targets) => targets.map((target) => {
      if (target.columnNames.length === 0) {
        return {
          targetId: target.targetId,
          tableName: target.tableName,
          rows: [],
          representativeValues: [],
        };
      }
      const columnName = target.columnNames[0]!;
      return {
        targetId: target.targetId,
        tableName: target.tableName,
        rows: wardValues.map((value) => ({ fields: [{ name: columnName, value }] })),
        representativeValues: [{ column: columnName, values: wardValues }],
      };
    })),
  };
  const completer: ModelCompleter = {
    complete: vi.fn(async (request) => {
      const context = JSON.parse(request.messages[1]!.content.split("\n").slice(1).join("\n"));
      return JSON.stringify({
        results: context.targets.map((target: { targetId: string }) => ({
          targetId: target.targetId,
          outcome: "generated",
          description: "Descrizione generata.",
        })),
      });
    }),
  };
  const models: MetadataGenerationModels = {
    catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
    resolve: () => ({
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    }),
  };
  const worker = new DescriptionGenerationWorker(
    repository,
    {
      read: vi.fn(async () => ({
        workspace: { workspace: { language: "it" } },
        revision: {},
      })),
    } as unknown as WorkspaceRegistry,
    models,
    completer,
    new CatalogOperationCoordinator(),
    sourceSampler,
  );

  const run = await worker.start(
    database.id,
    "openai-mini",
    "selected_columns",
    [email.id, phone.id, ward.id],
  );
  await worker.waitForRun(run.id);

  expect(sourceSampler.sample).toHaveBeenCalledWith(
    expect.objectContaining({ id: database.id }),
    [
      { targetId: email.id, tableName: table.name, columnNames: [] },
      { targetId: phone.id, tableName: table.name, columnNames: [] },
      { targetId: ward.id, tableName: table.name, columnNames: [ward.name] },
    ],
    expect.any(AbortSignal),
  );
  const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
  const context = JSON.parse(request.messages[1]!.content.split("\n").slice(1).join("\n"));
  const targets = new Map(
    context.targets.map((target: { targetId: string }) => [target.targetId, target]),
  );
  expect(targets.get(email.id)).toMatchObject({
    sourceSample: {
      rows: expect.arrayContaining([
        { fields: [{ name: email.name, value: "marta.rossi@example.com" }] },
      ]),
      representativeValues: [{
        column: email.name,
        values: expect.arrayContaining(["marta.rossi@example.com"]),
      }],
    },
  });
  expect(targets.get(phone.id)).toMatchObject({
    sourceSample: {
      rows: expect.arrayContaining([
        { fields: [{ name: phone.name, value: "+39 02 5550 1001" }] },
      ]),
      representativeValues: [{
        column: phone.name,
        values: expect.arrayContaining(["+39 02 5550 1001"]),
      }],
    },
  });
  expect(targets.get(ward.id)).toMatchObject({
    sourceSample: {
      rows: wardValues.map((value) => ({ fields: [{ name: ward.name, value }] })),
      representativeValues: [{ column: ward.name, values: wardValues }],
    },
  });
});

test("continues metadata-only with one safe warning when source sampling is unavailable", async () => {
  const repository = new MemoryCatalogRepository();
  const database = await repository.create({
    workspaceId: "psd-clinical",
    engine: "postgres",
    databaseName: "warehouse",
    schema: "datawarehouse",
    binding: {
      transport: "rest_api",
      baseUrl: "https://dwh.example.test",
      restPath: "/rpc/run_query",
      restAuth: "none",
    },
  });
  await repository.applySchemaSync(database.id, database.version, "all", [], {
    schemaVersion: 1,
    capabilities: { tables: "available", columns: "available", relationships: "available" },
    tables: [{ name: "patients", sourceComment: null }],
    columns: [{
      tableName: "patients",
      name: "status",
      ordinalPosition: 1,
      dataType: "text",
      isNullable: true,
      defaultExpression: null,
      primaryKeyPosition: null,
      sourceComment: null,
    }],
    relationships: [],
  });
  const table = (await repository.listTables(database.id))[0]!;
  const column = (await repository.listColumns(database.id, table.id))[0]!;
  const samplingFailureSecret = "UNAVAILABLE_SAMPLE_DETAIL_48b1f1";
  const sourceSampler: DescriptionSourceSampler = {
    sample: vi.fn(async () => { throw new Error(samplingFailureSecret); }),
  };
  const completer: ModelCompleter = {
    complete: vi.fn(async () => JSON.stringify({
      results: [{
        targetId: column.id,
        outcome: "generated",
        description: "Stato del paziente.",
      }],
    })),
  };
  const models: MetadataGenerationModels = {
    catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
    resolve: () => ({
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    }),
  };
  const worker = new DescriptionGenerationWorker(
    repository,
    {
      read: vi.fn(async () => ({
        workspace: { workspace: { language: "it" } },
        revision: {},
      })),
    } as unknown as WorkspaceRegistry,
    models,
    completer,
    new CatalogOperationCoordinator(),
    sourceSampler,
  );

  const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
  await worker.waitForRun(run.id);

  expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
    status: "completed",
    processed: 1,
    generated: 1,
    failed: 0,
  });
  const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
  expect(request.messages[1]?.content).not.toMatch(/sourceSample|UNAVAILABLE_SAMPLE_DETAIL/);
  const events = await repository.listDescriptionGenerationEvents(run.id);
  expect(events.filter((event) => event.level === "warning")).toEqual([
    expect.objectContaining({
      message: "Source samples unavailable for this batch; generation continued with catalog metadata only.",
    }),
  ]);
  expect(JSON.stringify(events)).not.toContain(samplingFailureSecret);
});