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

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

// ── Reranker with bounded scoring ──────────────────────────────────────────────

/**
 * Rerank chunks using a bounded scoring function.
 *
 * OPTIMIZATIONS over the original:
 *  1. Pre-computes query n-grams once instead of per-chunk.
 *  2. Uses a flat Set for token lookup (was already a Set, but now
 *     query token set is computed once via getTokenSet cache).
 *  3. N-gram overlap uses a local Set built from the chunk text only
 *     when needed (short-circuits on 0-length input).
 *  4. Source name boost uses Set.has() instead of String.includes() loop
 *     for O(1) per-token lookup.
 */

function getNGrams(text: string, n: number = 3): Set<string> {
  const clean = normalizeWhitespace(text.substring(0, 150)).toLowerCase();
  if (clean.length < n) return new Set([clean]);
  const grams = new Set<string>();
  for (let i = 0; i <= clean.length - n; i++) {
    grams.add(clean.substring(i, i + n));
  }
  return grams;
}

function ngramOverlap(queryGrams: Set<string>, text: string): number {
  const clean = normalizeWhitespace(text.substring(0, 150)).toLowerCase();
  if (clean.length < 3 || queryGrams.size === 0) return 0;

  const textGrams = new Set<string>();
  for (let i = 0; i <= clean.length - 3; i++) {
    textGrams.add(clean.substring(i, i + 3));
  }

  let matches = 0;
  for (const gram of queryGrams) {
    if (textGrams.has(gram)) matches++;
  }
  return matches / queryGrams.size;
}

export function rerankChunks(chunks: RankedChunk[], query: string): RankedChunk[] {
  if (chunks.length <= 2) return chunks.slice();

  // Pre-compute query tokens and n-grams once.
  const queryTokens = getTokenSet(query);
  const queryGrams = getNGrams(query);

  // Pre-build a Set from query tokens for O(1) source name lookup.
  const queryTokenSet = new Set(queryTokens);

  const scored = chunks.map((chunk) => {
    const textTokens = getTokenSet(chunk.text);

    // Lexical overlap: fraction of query tokens found in chunk text.
    const lexScore =
      queryTokenSet.size === 0
        ? 0
        : [...queryTokenSet].reduce(
            (c, t) => c + (textTokens.has(t) ? 1 : 0),
            0,
          ) / queryTokenSet.size;

    // N-gram overlap: fuzzy matching score.
    const ngramScore = ngramOverlap(queryGrams, chunk.text);

    // Source name boost: O(1) per token via Set.has().
    let boost = 0;
    const src = chunk.sourceName;
    if (src && queryTokenSet.size > 0) {
      const srcLower = normalizeWhitespace(src.toLowerCase());
      for (const token of queryTokenSet) {
        if (srcLower.includes(token)) {
          boost = 0.03;
          break;
        }
      }
    }

    // Bounded score: clamp base to [0, 1], add bounded bonuses.
    const base = Math.max(0, Math.min(1, chunk.score));
    const totalBonus = Math.min(0.5, lexScore * 0.3 + ngramScore * 0.4 + boost);

    return {
      ...chunk,
      score: Math.min(1.5, base + totalBonus),
      confidence: Math.max(0, Math.min(1, base * 0.7 + ngramScore * 0.3)),
    };
  });

  scored.sort((a, b) => {
    if (b.score !== a.score) return b.score - a.score;
    const cA = a.confidence ?? 0;
    const cB = b.confidence ?? 0;
    if (cB !== cA) return cB - cA;
    return a.text.length - b.text.length;
  });

  return scored;
}

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

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

// ── Reranker with bounded scoring ──────────────────────────────────────────────

/**
 * Rerank chunks using a bounded scoring function.
 *
 * OPTIMIZATIONS over the original:
 *  1. Pre-computes query n-grams once instead of per-chunk.
 *  2. Uses a flat Set for token lookup (was already a Set, but now
 *     query token set is computed once via getTokenSet cache).
 *  3. N-gram overlap uses a local Set built from the chunk text only
 *     when needed (short-circuits on 0-length input).
 *  4. Source name boost uses Set.has() instead of String.includes() loop
 *     for O(1) per-token lookup.
 */

function getNGrams(text: string, n: number = 3): Set<string> {
  const clean = normalizeWhitespace(text.substring(0, 150)).toLowerCase();
  if (clean.length < n) return new Set([clean]);
  const grams = new Set<string>();
  for (let i = 0; i <= clean.length - n; i++) {
    grams.add(clean.substring(i, i + n));
  }
  return grams;
}

function ngramOverlap(queryGrams: Set<string>, text: string): number {
  const clean = normalizeWhitespace(text.substring(0, 150)).toLowerCase();
  if (clean.length < 3 || queryGrams.size === 0) return 0;

  const textGrams = new Set<string>();
  for (let i = 0; i <= clean.length - 3; i++) {
    textGrams.add(clean.substring(i, i + 3));
  }

  let matches = 0;
  for (const gram of queryGrams) {
    if (textGrams.has(gram)) matches++;
  }
  return matches / queryGrams.size;
}

export function rerankChunks(chunks: RankedChunk[], query: string): RankedChunk[] {
  if (chunks.length <= 2) return chunks.slice();

  // Pre-compute query tokens and n-grams once.
  const queryTokens = getTokenSet(query);
  const queryGrams = getNGrams(query);

  // Pre-build a Set from query tokens for O(1) source name lookup.
  const queryTokenSet = new Set(queryTokens);

  const scored = chunks.map((chunk) => {
    const textTokens = getTokenSet(chunk.text);

    // Lexical overlap: fraction of query tokens found in chunk text.
    const lexScore =
      queryTokenSet.size === 0
        ? 0
        : [...queryTokenSet].reduce(
            (c, t) => c + (textTokens.has(t) ? 1 : 0),
            0,
          ) / queryTokenSet.size;

    // N-gram overlap: fuzzy matching score.
    const ngramScore = ngramOverlap(queryGrams, chunk.text);

    // Source name boost: O(1) per token via Set.has().
    let boost = 0;
    const src = chunk.sourceName;
    if (src && queryTokenSet.size > 0) {
      const srcLower = normalizeWhitespace(src.toLowerCase());
      for (const token of queryTokenSet) {
        if (srcLower.includes(token)) {
          boost = 0.03;
          break;
        }
      }
    }

    // Bounded score: clamp base to [0, 1], add bounded bonuses.
    const base = Math.max(0, Math.min(1, chunk.score));
    const totalBonus = Math.min(0.5, lexScore * 0.3 + ngramScore * 0.4 + boost);

    return {
      ...chunk,
      score: Math.min(1.5, base + totalBonus),
      confidence: Math.max(0, Math.min(1, base * 0.7 + ngramScore * 0.3)),
    };
  });

  scored.sort((a, b) => {
    if (b.score !== a.score) return b.score - a.score;
    const cA = a.confidence ?? 0;
    const cB = b.confidence ?? 0;
    if (cB !== cA) return cB - cA;
    return a.text.length - b.text.length;
  });

  return scored;
}