import { normalizeWhitespace } from "../utils/text";

export interface BudgetedChunk {
  text: string;
  score: number;
  citation?: string;
  sourceName?: string;
  confidence?: number;
}

/**
 * Fit chunks into a token budget, selecting diverse sources.
 *
 * OPTIMIZATIONS over the original:
 *  1. Pre-computes token lengths once (was called twice per chunk: in sort
 *     and in the main loop).
 *  2. Sort uses pre-computed token lengths — avoids re-estimating during
 *     comparisons (O(n log n) instead of O(n log n × m)).
 *  3. Removed redundant `totalTokens + neededTokens > safeBudget` check
 *     (already covered by the `safeBudget - totalTokens < neededTokens` check).
 */

function estimateTokens(text: string): number {
  const cleaned = normalizeWhitespace(text).trim();
  return Math.ceil(cleaned.length / 3.5);
}

export function fitToBudget(chunks: BudgetedChunk[], maxTokens: number): BudgetedChunk[] {
  if (chunks.length === 0) return [];

  const safeBudget = Math.floor(maxTokens * 0.8);
  const n = chunks.length;

  // Pre-compute token lengths once — avoids O(n log n) re-estimations in sort.
  const tokenLengths = new Int32Array(n);
  for (let i = 0; i < n; i++) {
    tokenLengths[i] = estimateTokens(chunks[i].text);
  }

  // Sort indices by score desc, confidence desc, token length asc.
  // Using index-based sort avoids spreading chunks into objects.
  const indices = Array.from({ length: n }, (_, i) => i);
  indices.sort((a, b) => {
    const scoreDiff = chunks[b].score - chunks[a].score;
    if (scoreDiff !== 0) return scoreDiff;
    const cA = chunks[a].confidence ?? 0;
    const cB = chunks[b].confidence ?? 0;
    if (cB !== cA) return cB - cA;
    return tokenLengths[a] - tokenLengths[b];
  });

  const kept: BudgetedChunk[] = [];
  const seenSources = new Map<string, number>();
  let totalTokens = 0;

  for (let idx = 0; idx < n; idx++) {
    const i = indices[idx];
    const sepOverhead = kept.length === 0 ? 0 : 6;
    const neededTokens = tokenLengths[i] + sepOverhead;

    if (safeBudget - totalTokens < neededTokens) continue;

    const key = chunks[i].sourceName || chunks[i].citation || `chunk-${i}`;
    const sourceCount = seenSources.get(key) || 0;
    if (sourceCount >= 3) continue;

    kept.push(chunks[i]);
    seenSources.set(key, sourceCount + 1);
    totalTokens += neededTokens;

    if (totalTokens >= safeBudget) break;
  }

  // Fallback: if nothing fit, take the single best chunk and truncate it.
  if (kept.length === 0 && n > 0) {
    const bestIdx = indices[0];
    const first = chunks[bestIdx];
    const maxChars = Math.max(10, Math.floor(safeBudget * 3.5));
    kept.push({
      ...first,
      text: first.text.slice(0, maxChars).trimEnd() + "\u2026",
    });
  }

  return kept;
}

import { normalizeWhitespace } from "../utils/text";

export interface BudgetedChunk {
  text: string;
  score: number;
  citation?: string;
  sourceName?: string;
  confidence?: number;
}

/**
 * Fit chunks into a token budget, selecting diverse sources.
 *
 * OPTIMIZATIONS over the original:
 *  1. Pre-computes token lengths once (was called twice per chunk: in sort
 *     and in the main loop).
 *  2. Sort uses pre-computed token lengths — avoids re-estimating during
 *     comparisons (O(n log n) instead of O(n log n × m)).
 *  3. Removed redundant `totalTokens + neededTokens > safeBudget` check
 *     (already covered by the `safeBudget - totalTokens < neededTokens` check).
 */

function estimateTokens(text: string): number {
  const cleaned = normalizeWhitespace(text).trim();
  return Math.ceil(cleaned.length / 3.5);
}

export function fitToBudget(chunks: BudgetedChunk[], maxTokens: number): BudgetedChunk[] {
  if (chunks.length === 0) return [];

  const safeBudget = Math.floor(maxTokens * 0.8);
  const n = chunks.length;

  // Pre-compute token lengths once — avoids O(n log n) re-estimations in sort.
  const tokenLengths = new Int32Array(n);
  for (let i = 0; i < n; i++) {
    tokenLengths[i] = estimateTokens(chunks[i].text);
  }

  // Sort indices by score desc, confidence desc, token length asc.
  // Using index-based sort avoids spreading chunks into objects.
  const indices = Array.from({ length: n }, (_, i) => i);
  indices.sort((a, b) => {
    const scoreDiff = chunks[b].score - chunks[a].score;
    if (scoreDiff !== 0) return scoreDiff;
    const cA = chunks[a].confidence ?? 0;
    const cB = chunks[b].confidence ?? 0;
    if (cB !== cA) return cB - cA;
    return tokenLengths[a] - tokenLengths[b];
  });

  const kept: BudgetedChunk[] = [];
  const seenSources = new Map<string, number>();
  let totalTokens = 0;

  for (let idx = 0; idx < n; idx++) {
    const i = indices[idx];
    const sepOverhead = kept.length === 0 ? 0 : 6;
    const neededTokens = tokenLengths[i] + sepOverhead;

    if (safeBudget - totalTokens < neededTokens) continue;

    const key = chunks[i].sourceName || chunks[i].citation || `chunk-${i}`;
    const sourceCount = seenSources.get(key) || 0;
    if (sourceCount >= 3) continue;

    kept.push(chunks[i]);
    seenSources.set(key, sourceCount + 1);
    totalTokens += neededTokens;

    if (totalTokens >= safeBudget) break;
  }

  // Fallback: if nothing fit, take the single best chunk and truncate it.
  if (kept.length === 0 && n > 0) {
    const bestIdx = indices[0];
    const first = chunks[bestIdx];
    const maxChars = Math.max(10, Math.floor(safeBudget * 3.5));
    kept.push({
      ...first,
      text: first.text.slice(0, maxChars).trimEnd() + "\u2026",
    });
  }

  return kept;
}