find-image

/**
 * Selects a bounded batch without letting the first indexed source consume
 * the entire embedding budget. Source order and per-source candidate order
 * remain stable.
 */
export function selectEmbeddingCandidates<T>(
  candidates: readonly T[],
  maxCandidates: number,
  sourceKey: (candidate: T) => string | undefined,
): T[] {
  const budget = Math.max(0, Math.floor(maxCandidates));
  if (budget === 0 || candidates.length === 0) return [];
  if (budget >= candidates.length) return [...candidates];

  const candidatesBySource = new Map<string, T[]>();
  for (const candidate of candidates) {
    const source = sourceKey(candidate) || "unknown";
    const sourceCandidates = candidatesBySource.get(source);
    if (sourceCandidates) {
      sourceCandidates.push(candidate);
    } else {
      candidatesBySource.set(source, [candidate]);
    }
  }

  const selected: T[] = [];
  while (selected.length < budget) {
    let selectedInRound = false;
    for (const sourceCandidates of candidatesBySource.values()) {
      const candidate = sourceCandidates.shift();
      if (!candidate) continue;
      selected.push(candidate);
      selectedInRound = true;
      if (selected.length === budget) break;
    }
    if (!selectedInRound) break;
  }

  return selected;
}