import type { RegulatoryChunk, RegulatoryRetrievalCandidate } from '../../src/compliance/regulatory-rag';

/**
 * In-memory BM25 over the chunks of one indexed corpus, the keyword half of the hybrid search (Phase 13F, user
 * decision 2026-10-09). The corpora are small (some hundred to some thousand chunks), so the index is built once per
 * corpus version from Qdrant and kept in memory; nothing is stored. Tokens are lower-cased words with a light German
 * and English suffix stemming, so `Berechtigungen` and `Berechtigung` meet.
 */
const STOP_WORDS = new Set(('der die das und oder ein eine einer eines zu zur zum von mit für auf in im ist sind werden wird sollte '
  + 'muss kann nicht den dem des als bei auch über durch wie the and or of to a an in on for is are be with this that by as it its '
  + 'from not should must can').split(' '));

export function lexicalTokens(text: string): string[] {
  return (text.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [])
    .filter(word => word.length > 2 && !STOP_WORDS.has(word))
    .map(word => word.replace(/(ungen|ung|en|er|es|e|n|s)$/u, ''));
}

export class RegulatoryLexicalIndex {
  private readonly documents: Array<Map<string, number>>;
  private readonly lengths: number[];
  private readonly documentFrequency = new Map<string, number>();
  private readonly averageLength: number;

  constructor(private readonly chunks: readonly RegulatoryChunk[]) {
    this.documents = chunks.map(chunk => {
      const counts = new Map<string, number>();
      for (const token of lexicalTokens(`${chunk.title} ${chunk.text}`)) counts.set(token, (counts.get(token) ?? 0) + 1);
      return counts;
    });
    this.lengths = this.documents.map(counts => [...counts.values()].reduce((total, count) => total + count, 0));
    for (const counts of this.documents) for (const token of counts.keys()) this.documentFrequency.set(token, (this.documentFrequency.get(token) ?? 0) + 1);
    this.averageLength = this.lengths.reduce((total, length) => total + length, 0) / Math.max(1, this.lengths.length);
  }

  get size(): number { return this.chunks.length; }

  /** Chunks with at least one query term, best first; equal scores keep corpus order, so results are deterministic. */
  search(query: string, limit: number): RegulatoryRetrievalCandidate[] {
    const terms = [...new Set(lexicalTokens(query))];
    const total = this.documents.length;
    const hits: RegulatoryRetrievalCandidate[] = [];
    this.documents.forEach((counts, index) => {
      let score = 0;
      for (const term of terms) {
        const frequency = counts.get(term);
        if (!frequency) continue;
        const documents = this.documentFrequency.get(term)!;
        const idf = Math.log(1 + (total - documents + 0.5) / (documents + 0.5));
        score += idf * (frequency * 2.2) / (frequency + 1.2 * (0.25 + 0.75 * this.lengths[index] / this.averageLength));
      }
      if (score > 0) hits.push({ score, chunk: this.chunks[index] });
    });
    return hits.sort((left, right) => right.score - left.score).slice(0, limit);
  }
}

export interface FusedControl {
  controlId: string;
  /** Reciprocal rank fusion score over the ranked lists; it orders candidates and is no relevance probability. */
  score: number;
  /** Every distinct chunk of the control any list returned, in list order. */
  chunks: RegulatoryChunk[];
}

/**
 * Reciprocal rank fusion (k = 60) of ranked chunk lists at control level: a control ranks in each list by its best
 * chunk. Ties fall back to the control id, so identical input yields identical candidates.
 */
export function fuseControlRankings(lists: ReadonlyArray<readonly RegulatoryRetrievalCandidate[]>, exclude?: (controlId: string) => boolean): FusedControl[] {
  const fused = new Map<string, { score: number; chunks: Map<string, RegulatoryChunk> }>();
  for (const list of lists) {
    const order: string[] = [];
    for (const hit of list) {
      const controlId = hit.chunk.controlId;
      if (exclude?.(controlId)) continue;
      const entry = fused.get(controlId) ?? { score: 0, chunks: new Map<string, RegulatoryChunk>() };
      if (!entry.chunks.has(hit.chunk.chunkId)) entry.chunks.set(hit.chunk.chunkId, hit.chunk);
      fused.set(controlId, entry);
      if (!order.includes(controlId)) order.push(controlId);
    }
    order.forEach((controlId, rank) => { fused.get(controlId)!.score += 1 / (60 + rank + 1); });
  }
  return [...fused.entries()]
    .map(([controlId, entry]) => ({ controlId, score: entry.score, chunks: [...entry.chunks.values()] }))
    .sort((left, right) => right.score - left.score || left.controlId.localeCompare(right.controlId, 'en', { numeric: true }));
}