import { describe, expect, it, vi } from 'vitest';
import type { RegulatoryControlDocument, RegulatoryRetrievalCandidate } from '../../src/domain/regulatory-rag';
import type { TextEmbeddingClient } from '../../server/regulatory/embedding-client';
import type { IndexedRegulatoryChunk, RegulatoryChunkIndex } from '../../server/regulatory/qdrant-regulatory-index';
import {
  DenseRegulatoryCandidateRetriever,
  RegulatoryIngestionService,
  RegulatoryRetrievalPipeline
} from '../../server/regulatory/regulatory-rag-service';

const document: RegulatoryControlDocument = {
  schemaVersion: 1, corpusId: 'nis2', corpusVersion: '2024-1', documentId: 'nis2-21-2-a-de',
  framework: 'NIS2', frameworkVersion: '2022/2555', controlId: '21(2)(a)', language: 'de',
  title: 'Risikoanalyse', text: 'Synthetischer Kontrolltext. '.repeat(30)
};

describe('regulatory RAG services', () => {
  it('embeds every chunk before replacing documents in the derived index', async () => {
    const embeddingBatches: string[][] = [];
    const embeddings: TextEmbeddingClient = {
      model: 'multilingual-model', dimensions: 2,
      embed: async texts => {
        embeddingBatches.push(texts);
        return texts.map((_, index) => [index, 1]);
      }
    };
    const ensureCollection = vi.fn(async () => {});
    const replaceDocuments = vi.fn(async (_chunks: IndexedRegulatoryChunk[]) => {});
    const index: RegulatoryChunkIndex = { ensureCollection, replaceDocuments, searchDense: vi.fn() };
    const service = new RegulatoryIngestionService(embeddings, index, { maximumChunkCharacters: 220, embeddingBatchSize: 2 });

    const result = await service.ingest([document]);

    expect(result.documentCount).toBe(1);
    expect(result.chunkCount).toBeGreaterThan(1);
    expect(embeddingBatches.every(batch => batch.length <= 2)).toBe(true);
    expect(ensureCollection).toHaveBeenCalledOnce();
    expect(replaceDocuments).toHaveBeenCalledOnce();
    expect(replaceDocuments.mock.calls[0][0]).toHaveLength(result.chunkCount);
  });

  it('supports multilingual dense candidates and an optional later reranker stage', async () => {
    const embed = vi.fn(async () => [[0.5, 0.5]]);
    const candidates: RegulatoryRetrievalCandidate[] = Array.from({ length: 8 }, (_, index) => ({
      score: 1 - index / 10,
      chunk: { ...document, chunkId: `chunk-${index}`, chunkIndex: index, textSha256: String(index).padStart(64, '0') }
    }));
    const searchDense = vi.fn(async () => candidates);
    const dense = new DenseRegulatoryCandidateRetriever(
      { model: 'multilingual-model', dimensions: 2, embed },
      { ensureCollection: vi.fn(), replaceDocuments: vi.fn(), searchDense }
    );
    const rerank = vi.fn(async (_query: string, values: RegulatoryRetrievalCandidate[]) => values.slice().reverse());
    const retrieval = new RegulatoryRetrievalPipeline(dense, { rerank }, 4);

    const result = await retrieval.retrieve({ query: 'How is the German control relevant?', limit: 2, filters: { language: 'de' } });

    expect(embed).toHaveBeenCalledWith(['How is the German control relevant?']);
    expect(searchDense).toHaveBeenCalledWith([0.5, 0.5], 8, { language: 'de' });
    expect(rerank).toHaveBeenCalledOnce();
    expect(result.map(item => item.chunk.chunkId)).toEqual(['chunk-7', 'chunk-6']);
  });
});