import type {
  RegulatoryChunk,
  RegulatoryControlDocument,
  RegulatoryReranker,
  RegulatoryRetrievalCandidate,
  RegulatoryRetrievalRequest,
  RegulatoryRetriever
} from '../../src/domain/regulatory-rag';
import type { TextEmbeddingClient } from './embedding-client';
import { chunkRegulatoryControl, regulatoryChunkEmbeddingText, validateRegulatoryControlDocument } from './regulatory-chunker';
import type { RegulatoryChunkIndex, RegulatorySparseVector } from './qdrant-regulatory-index';

const HYBRID_RRF_K = 60;
const BM25_K1 = 1.2;
const BM25_B = 0.75;

export interface RegulatoryCandidateRetriever {
  retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]>;
}

export class DenseRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
  constructor(
    private readonly embeddings: TextEmbeddingClient,
    private readonly index: RegulatoryChunkIndex
  ) {}

  async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
    const [vector] = await this.embeddings.embed([validatedQuery(request.query)]);
    return this.index.searchDense(vector, validatedLimit(request.limit), request.filters);
  }
}

export class SparseRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
  constructor(private readonly index: RegulatoryChunkIndex) {}

  async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
    const vector = regulatoryQuerySparseVector(validatedQuery(request.query));
    if (!vector) return [];
    return this.index.searchSparse(vector, validatedLimit(request.limit), request.filters);
  }
}

export class HybridRegulatoryCandidateRetriever implements RegulatoryCandidateRetriever {
  constructor(
    private readonly dense: RegulatoryCandidateRetriever,
    private readonly sparse: RegulatoryCandidateRetriever
  ) {}

  async retrieveCandidates(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
    const limit = validatedLimit(request.limit);
    const [dense, sparse] = await Promise.all([
      this.dense.retrieveCandidates({ ...request, limit }),
      this.sparse.retrieveCandidates({ ...request, limit })
    ]);
    return reciprocalRankFuseCandidates([dense, sparse], limit);
  }
}

export class RegulatoryRetrievalPipeline implements RegulatoryRetriever {
  constructor(
    private readonly candidates: RegulatoryCandidateRetriever,
    private readonly reranker?: RegulatoryReranker,
    private readonly rerankCandidateMultiplier = 4
  ) {
    if (!Number.isInteger(rerankCandidateMultiplier) || rerankCandidateMultiplier < 1 || rerankCandidateMultiplier > 10) {
      throw new Error('rerankCandidateMultiplier must be an integer between 1 and 10.');
    }
  }

  async retrieve(request: RegulatoryRetrievalRequest): Promise<RegulatoryRetrievalCandidate[]> {
    const limit = validatedLimit(request.limit);
    const query = validatedQuery(request.query);
    const candidateLimit = this.reranker ? Math.min(100, limit * this.rerankCandidateMultiplier) : limit;
    const candidates = await this.candidates.retrieveCandidates({ ...request, query, limit: candidateLimit });
    if (!this.reranker) return candidates.slice(0, limit);
    return (await this.reranker.rerank(query, candidates)).slice(0, limit);
  }
}

export interface RegulatoryIngestionOptions {
  maximumChunkCharacters?: number;
  embeddingBatchSize?: number;
}

export interface RegulatoryIngestionResult {
  documentCount: number;
  chunkCount: number;
  embeddingModel: string;
}

export class RegulatoryIngestionService {
  private readonly maximumChunkCharacters: number;
  private readonly embeddingBatchSize: number;

  constructor(
    private readonly embeddings: TextEmbeddingClient,
    private readonly index: RegulatoryChunkIndex,
    options: RegulatoryIngestionOptions = {}
  ) {
    this.maximumChunkCharacters = options.maximumChunkCharacters ?? 1_600;
    this.embeddingBatchSize = options.embeddingBatchSize ?? 32;
    if (!Number.isInteger(this.embeddingBatchSize) || this.embeddingBatchSize < 1 || this.embeddingBatchSize > 128) {
      throw new Error('embeddingBatchSize must be an integer between 1 and 128.');
    }
  }

  async ingest(documents: RegulatoryControlDocument[], onProgress?: (completedChunks: number, totalChunks: number) => Promise<void> | void): Promise<RegulatoryIngestionResult> {
    if (!documents.length) throw new Error('At least one regulatory control document is required.');
    const documentKeys = new Set<string>();
    const chunks = documents.flatMap(input => {
      const document = validateRegulatoryControlDocument(input);
      const key = `${document.corpusId}\0${document.documentId}`;
      if (documentKeys.has(key)) throw new Error(`Duplicate regulatory document identity: ${document.documentId}.`);
      documentKeys.add(key);
      return chunkRegulatoryControl(document, { maximumCharacters: this.maximumChunkCharacters });
    });
    const vectors: number[][] = [];
    await onProgress?.(0, chunks.length);
    for (let offset = 0; offset < chunks.length; offset += this.embeddingBatchSize) {
      const batch = chunks.slice(offset, offset + this.embeddingBatchSize).map(regulatoryChunkEmbeddingText);
      vectors.push(...await this.embeddings.embed(batch));
      await onProgress?.(Math.min(offset + batch.length, chunks.length), chunks.length);
    }
    if (vectors.length !== chunks.length) throw new Error('Embedding provider returned an unexpected vector count.');
    const sparseVectors = buildRegulatorySparseVectors(chunks);
    await this.index.ensureCollection();
    await this.index.replaceDocuments(chunks.map((chunk, index) => ({
      chunk,
      embedding: vectors[index],
      ...(sparseVectors.get(chunk.chunkId) ? { sparseVector: sparseVectors.get(chunk.chunkId) } : {})
    })));
    return { documentCount: documents.length, chunkCount: chunks.length, embeddingModel: this.embeddings.model };
  }
}

/** Builds deterministic hashed BM25 document vectors without a scanner-specific vocabulary. */
export function buildRegulatorySparseVectors(chunks: RegulatoryChunk[]): Map<string, RegulatorySparseVector> {
  const documents = chunks.map(chunk => tokenFrequencies(`${chunk.title} ${chunk.title} ${chunk.text}`));
  const documentFrequency = new Map<number, number>();
  for (const document of documents) {
    for (const index of document.frequencies.keys()) documentFrequency.set(index, (documentFrequency.get(index) ?? 0) + 1);
  }
  const averageLength = documents.reduce((sum, document) => sum + document.length, 0) / Math.max(documents.length, 1);
  const vectors = new Map<string, RegulatorySparseVector>();
  documents.forEach((document, documentIndex) => {
    if (!document.frequencies.size) return;
    const entries = [...document.frequencies.entries()].map(([index, frequency]) => {
      const matchingDocuments = documentFrequency.get(index) ?? 1;
      const inverseDocumentFrequency = Math.log(1 + (documents.length - matchingDocuments + 0.5) / (matchingDocuments + 0.5));
      const normalizedFrequency = frequency * (BM25_K1 + 1)
        / (frequency + BM25_K1 * (1 - BM25_B + BM25_B * document.length / Math.max(averageLength, 1)));
      return [index, inverseDocumentFrequency * normalizedFrequency] as const;
    }).sort(([left], [right]) => left - right);
    vectors.set(chunks[documentIndex].chunkId, {
      indices: entries.map(([index]) => index),
      values: entries.map(([, score]) => score)
    });
  });
  return vectors;
}

export function regulatoryQuerySparseVector(query: string): RegulatorySparseVector | undefined {
  const indices = [...new Set(regulatoryTokens(query).map(sparseIndex))].sort((left, right) => left - right);
  return indices.length ? { indices, values: indices.map(() => 1) } : undefined;
}

export function reciprocalRankFuseCandidates(
  rankings: RegulatoryRetrievalCandidate[][],
  limit: number
): RegulatoryRetrievalCandidate[] {
  const fused = new Map<string, { score: number; bestRank: number; candidate: RegulatoryRetrievalCandidate }>();
  for (const ranking of rankings) {
    const seen = new Set<string>();
    let rank = 0;
    for (const candidate of ranking) {
      const controlId = candidate.chunk.controlId;
      if (seen.has(controlId)) continue;
      seen.add(controlId);
      rank += 1;
      const existing = fused.get(controlId);
      const entry = existing ?? { score: 0, bestRank: Number.MAX_SAFE_INTEGER, candidate };
      entry.score += 1 / (HYBRID_RRF_K + rank);
      if (rank < entry.bestRank) {
        entry.bestRank = rank;
        entry.candidate = candidate;
      }
      fused.set(controlId, entry);
    }
  }
  return [...fused.values()]
    .sort((left, right) => right.score - left.score
      || left.bestRank - right.bestRank
      || left.candidate.chunk.controlId.localeCompare(right.candidate.chunk.controlId, 'en', { numeric: true }))
    .slice(0, limit)
    .map(entry => ({ ...entry.candidate, score: entry.score }));
}

function tokenFrequencies(value: string): { frequencies: Map<number, number>; length: number } {
  const frequencies = new Map<number, number>();
  const tokens = regulatoryTokens(value);
  for (const token of tokens) {
    const index = sparseIndex(token);
    frequencies.set(index, (frequencies.get(index) ?? 0) + 1);
  }
  return { frequencies, length: tokens.length };
}

function regulatoryTokens(value: string): string[] {
  return (value.normalize('NFKD').toLocaleLowerCase('und').replace(/\p{M}/gu, '').match(/[\p{L}\p{N}]+/gu) ?? [])
    .filter(token => token.length > 1);
}

function sparseIndex(token: string): number {
  let hash = 0x811c9dc5;
  for (let index = 0; index < token.length; index += 1) {
    hash ^= token.charCodeAt(index);
    hash = Math.imul(hash, 0x01000193);
  }
  return hash >>> 0;
}

function validatedQuery(value: string): string {
  if (typeof value !== 'string' || !value.trim()) throw new Error('Regulatory retrieval query must be non-empty.');
  if (value.length > 20_000) throw new Error('Regulatory retrieval query exceeds 20000 characters.');
  return value.trim();
}

function validatedLimit(value: number): number {
  if (!Number.isInteger(value) || value < 1 || value > 100) throw new Error('Regulatory retrieval limit must be between 1 and 100.');
  return value;
}