import type { Scored, SourceRef } from './types';

/**
 * Reciprocal Rank Fusion (docs/AI-ASSISTANT-CONCEPT.md §7 step 3):
 * score(d) = Σ 1/(k + rank_i(d)), k = 60.
 *
 * RRF only reads rank, not the underlying score — so it needs no calibration
 * between BM25 (FTS) and cosine (embedding) scores, and degrades to a single
 * retriever with no code branch when one leg is absent (just pass an empty
 * array for that leg).
 */
const RRF_K = 60;

/**
 * Deliberately excludes collectionId: a JMAP email can live in more than one
 * mailbox, and the FTS leg and the embedding leg may legitimately report a
 * different "primary" one for the same message. itemId is already the real
 * identity within an account - including collectionId here would let the
 * same email be counted twice instead of properly fused.
 */
function refKey(ref: SourceRef): string {
  return `${ref.product}:${ref.accountId}:${ref.itemId}:${ref.chunkIx}`;
}

export function reciprocalRankFusion(
  legs: Scored<SourceRef>[][],
  limit: number,
): Scored<SourceRef>[] {
  const fused = new Map<string, { ref: SourceRef; score: number }>();

  for (const leg of legs) {
    leg.forEach((hit, index) => {
      const key = refKey(hit.ref);
      const rank = index + 1;
      const contribution = 1 / (RRF_K + rank);
      const existing = fused.get(key);
      if (existing) {
        existing.score += contribution;
      } else {
        fused.set(key, { ref: hit.ref, score: contribution });
      }
    });
  }

  return [...fused.values()]
    .sort((a, b) => b.score - a.score)
    .slice(0, limit);
}