/**
 * measure-retrieval.ts
 *
 * Pure retrieval quality measurement functions.
 *
 * Evaluates retrieval quality by comparing what Sanity's text search
 * returns against the manually-annotated canonical documents for each
 * evaluation task. Produces Recall@K and NDCG@K metrics.
 *
 * This answers: "Can a retriever find the docs an LLM actually needs?"
 *
 * Migrated from lib/measure-retrieval.ts — no process.argv/process.env,
 * accepts rootDir and retriever function as parameters.
 */
import { resolveMappings } from "./resolve-mappings.js";
// ---------------------------------------------------------------------------
// Metrics (pure, exported for testing)
// ---------------------------------------------------------------------------
export function calculateRecall(canonical, retrieved, k) {
    const retrievedSet = new Set(retrieved.slice(0, k));
    const hits = canonical.filter((doc) => retrievedSet.has(doc)).length;
    return canonical.length === 0 ? 0 : hits / canonical.length;
}
export function calculateNDCG(canonical, retrieved, k) {
    const canonicalSet = new Set(canonical);
    // Discounted Cumulative Gain
    let dcg = 0;
    for (let i = 0; i < Math.min(k, retrieved.length); i++) {
        if (canonicalSet.has(retrieved[i])) {
            dcg += 1 / Math.log2(i + 2); // +2 because log2(1) = 0
        }
    }
    // Ideal DCG
    let idcg = 0;
    for (let i = 0; i < Math.min(k, canonical.length); i++) {
        idcg += 1 / Math.log2(i + 2);
    }
    return idcg === 0 ? 0 : dcg / idcg;
}
/**
 * Run retrieval quality measurement for all tasks.
 *
 * @returns A full RetrievalSummary with per-task, per-area, and overall metrics.
 */
export async function measureRetrieval(options) {
    const { rootDir, retriever, onProgress } = options;
    const mappings = resolveMappings(rootDir);
    const results = [];
    for (const [area, areaData] of Object.entries(mappings.feature_areas)) {
        for (const task of areaData.tasks) {
            const canonicalSlugs = task.canonical_docs.map((d) => d.slug);
            // Use the task description as a search query
            const retrieved = await retriever(task.description, 10);
            const result = {
                canonical_docs: canonicalSlugs,
                feature_area: area,
                ndcg_at_10: calculateNDCG(canonicalSlugs, retrieved, 10),
                recall_at_5: calculateRecall(canonicalSlugs, retrieved, 5),
                recall_at_10: calculateRecall(canonicalSlugs, retrieved, 10),
                retrieved_docs: retrieved,
                task_id: task.id,
            };
            results.push(result);
            onProgress?.(area, task.id, result);
        }
    }
    // Aggregate by feature area
    const byArea = {};
    for (const area of Object.keys(mappings.feature_areas)) {
        const areaResults = results.filter((r) => r.feature_area === area);
        if (areaResults.length === 0)
            continue;
        byArea[area] = {
            avg_ndcg_at_10: areaResults.reduce((s, r) => s + r.ndcg_at_10, 0) / areaResults.length,
            avg_recall_at_5: areaResults.reduce((s, r) => s + r.recall_at_5, 0) / areaResults.length,
            avg_recall_at_10: areaResults.reduce((s, r) => s + r.recall_at_10, 0) /
                areaResults.length,
            task_count: areaResults.length,
        };
    }
    // Overall
    const overall = {
        avg_ndcg_at_10: results.reduce((s, r) => s + r.ndcg_at_10, 0) / (results.length || 1),
        avg_recall_at_5: results.reduce((s, r) => s + r.recall_at_5, 0) / (results.length || 1),
        avg_recall_at_10: results.reduce((s, r) => s + r.recall_at_10, 0) / (results.length || 1),
    };
    return { by_area: byArea, overall, results };
}
/**
 * Format a retrieval summary for console display.
 */
export function formatRetrievalTable(summary) {
    const lines = [];
    lines.push("=".repeat(70));
    lines.push("RETRIEVAL QUALITY SUMMARY");
    lines.push("=".repeat(70));
    lines.push("");
    lines.push("| Feature Area        | Recall@5 | Recall@10 | NDCG@10 | Tasks |");
    lines.push("|---------------------|----------|-----------|---------|-------|");
    for (const [area, stats] of Object.entries(summary.by_area)) {
        lines.push(`| ${area.padEnd(19)} | ${(stats.avg_recall_at_5 * 100).toFixed(1).padStart(7)}% | ` +
            `${(stats.avg_recall_at_10 * 100).toFixed(1).padStart(8)}% | ` +
            `${(stats.avg_ndcg_at_10 * 100).toFixed(1).padStart(6)}% | ` +
            `${stats.task_count.toString().padStart(5)} |`);
    }
    lines.push("");
    lines.push(`Overall:  Recall@5=${(summary.overall.avg_recall_at_5 * 100).toFixed(1)}%  ` +
        `Recall@10=${(summary.overall.avg_recall_at_10 * 100).toFixed(1)}%  ` +
        `NDCG@10=${(summary.overall.avg_ndcg_at_10 * 100).toFixed(1)}%`);
    return lines.join("\n");
}
