/**
 * Maximal Marginal Relevance (MMR) re-ranking — Phase 2.
 *
 * Ported from upstream `extensions/memory-core/src/memory/mmr.ts` with
 * minor adaptation for our `MemoryEntry` shape (content vs snippet, id vs
 * path+startLine). Behavior is identical: normalize scores, iteratively
 * pick the item that maximizes `λ * relevance - (1-λ) * max_similarity`
 * using Jaccard on tokenized content.
 *
 * @see Carbonell & Goldstein, "The Use of MMR, Diversity-Based Reranking" (1998)
 */

export interface MMRConfig {
  /** Opt-in. Upstream default is false. */
  enabled: boolean;
  /** 0 = max diversity, 1 = max relevance. Upstream default 0.7. */
  lambda: number;
}

export const DEFAULT_MMR_CONFIG: MMRConfig = {
  enabled: false,
  lambda: 0.7,
};

/**
 * CJK Unified Ideographs + Extension A, Hiragana/Katakana, Hangul.
 * These lack whitespace boundaries so we must tokenize them differently.
 */
const CJK_RE =
  /[\u3040-\u309f\u30a0-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uac00-\ud7af\u1100-\u11ff]/;

/**
 * Tokenize content into a set for Jaccard similarity.
 *
 * ASCII: alphanumeric + underscore runs, lowercased.
 * CJK: each char becomes a unigram; consecutive pairs become a bigram.
 * Non-adjacent CJK chars (e.g. `我a好`) do NOT form a bigram.
 */
export function tokenize(text: string): Set<string> {
  const lower = (text ?? '').toLowerCase();
  const ascii = lower.match(/[a-z0-9_]+/g) ?? [];

  const chars = Array.from(lower);
  const cjkData: Array<{ char: string; index: number }> = [];
  for (let i = 0; i < chars.length; i++) {
    if (CJK_RE.test(chars[i])) cjkData.push({ char: chars[i], index: i });
  }

  const bigrams: string[] = [];
  for (let i = 0; i < cjkData.length - 1; i++) {
    if (cjkData[i + 1].index === cjkData[i].index + 1) {
      bigrams.push(cjkData[i].char + cjkData[i + 1].char);
    }
  }

  return new Set([...ascii, ...bigrams, ...cjkData.map((d) => d.char)]);
}

export function jaccardSimilarity(a: Set<string>, b: Set<string>): number {
  if (a.size === 0 && b.size === 0) return 1;
  if (a.size === 0 || b.size === 0) return 0;

  const [smaller, larger] = a.size <= b.size ? [a, b] : [b, a];
  let intersection = 0;
  for (const t of smaller) if (larger.has(t)) intersection++;
  const union = a.size + b.size - intersection;
  return union === 0 ? 0 : intersection / union;
}

export function textSimilarity(a: string, b: string): number {
  return jaccardSimilarity(tokenize(a), tokenize(b));
}

export function computeMMRScore(
  relevance: number,
  maxSimilarity: number,
  lambda: number
): number {
  return lambda * relevance - (1 - lambda) * maxSimilarity;
}

export interface MMRItem {
  id: string;
  score: number;
  content: string;
}

/**
 * Re-rank items using MMR. Returns a new array in MMR order.
 */
export function mmrRerank<T extends MMRItem>(
  items: T[],
  config: Partial<MMRConfig> = {}
): T[] {
  const enabled = config.enabled ?? DEFAULT_MMR_CONFIG.enabled;
  const rawLambda = config.lambda ?? DEFAULT_MMR_CONFIG.lambda;

  if (!enabled || items.length <= 1) return [...items];

  const lambda = Math.max(0, Math.min(1, rawLambda));
  if (lambda === 1) return [...items].sort((a, b) => b.score - a.score);

  const tokenCache = new Map<string, Set<string>>();
  for (const item of items) tokenCache.set(item.id, tokenize(item.content));

  const scores = items.map((i) => i.score);
  const maxScore = Math.max(...scores);
  const minScore = Math.min(...scores);
  const range = maxScore - minScore;
  const normalize = (s: number): number =>
    range === 0 ? 1 : (s - minScore) / range;

  const selected: T[] = [];
  const remaining = new Set(items);

  while (remaining.size > 0) {
    let best: T | null = null;
    let bestMMR = -Infinity;
    for (const cand of remaining) {
      const rel = normalize(cand.score);
      const candTokens = tokenCache.get(cand.id)!;
      let maxSim = 0;
      for (const sel of selected) {
        const sim = jaccardSimilarity(candTokens, tokenCache.get(sel.id)!);
        if (sim > maxSim) maxSim = sim;
      }
      const mmr = computeMMRScore(rel, maxSim, lambda);
      if (
        mmr > bestMMR ||
        (mmr === bestMMR && cand.score > (best?.score ?? -Infinity))
      ) {
        bestMMR = mmr;
        best = cand;
      }
    }
    if (!best) break;
    selected.push(best);
    remaining.delete(best);
  }

  return selected;
}

/**
 * Adapter: apply MMR to an array of MemoryEntry-shaped hits.
 *
 * Uses (path|id|index) as the stable ID so two hits from the same file at
 * different content still get distinct MMR identities.
 */
export function applyMMRToMemoryHits<
  T extends { id: string; path: string; content: string; score: number },
>(results: T[], config: Partial<MMRConfig> = {}): T[] {
  if (results.length <= 1) return results;
  const byId = new Map<string, T>();
  const items: MMRItem[] = results.map((r, i) => {
    const id = `${r.path}#${r.id}#${i}`;
    byId.set(id, r);
    return { id, score: r.score, content: r.content };
  });
  return mmrRerank(items, config).map((m) => byId.get(m.id)!);
}