#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss,
    clippy::too_many_lines
)]
//! Profile-driven optimization harness for the ripvec engine.
//!
//! Drives the SAME end-to-end runtime path the MCP `search` tool
//! takes:
//!
//!     RipvecIndex::from_root
//!       (encoder + chunks + embeddings + BM25 + corpus_class)
//!       + repo_map::build_graph -> pagerank_lookup (PageRank ON)
//!     -> index.search(Hybrid, filter_paths=scope-derived)
//!     -> if should_rerank: CrossEncoderRerank layer applied
//!     -> score-threshold + top-k slice
//!
//! Why this exists (vs `semble_bench.rs`):
//!
//! - We're not benchmarking parity against semble any more. We're
//!   profiling and optimizing ripvec's own end-to-end. semble_bench
//!   intentionally pinned the path to `pagerank=None` for matched-
//!   model fidelity; that's the wrong fixture for `tracemeld:profile`
//!   because it leaves PageRank — a real, runtime-path component —
//!   out of the picture.
//!
//! - We need recall + precision tracked at every bench, not just
//!   NDCG. The optimization loop is "did this change move speed
//!   without regressing quality" — both numbers anchor the answer.
//!
//! Usage:
//!     cargo run --release --example corpus_bench -- \
//!         <corpus_root> <annotations.json> [--model REPO] \
//!         [--scope code|docs|all] [--no-rerank] [--rerank]
//!         [--repeats N]

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use ripvec_core::chunk::CodeChunk;
use ripvec_core::embed::{Scope, SearchConfig};
use ripvec_core::encoder::ripvec::dense::{DEFAULT_MODEL_REPO, StaticEncoder};
use ripvec_core::encoder::ripvec::index::RipvecIndex;
use ripvec_core::encoder::ripvec::ranking::is_symbol_query;
use ripvec_core::hybrid::{SearchMode, pagerank_lookup};
use ripvec_core::profile::Profiler;
use ripvec_core::ranking::{CrossEncoderRerank, RankingLayer, apply_chain};
use ripvec_core::repo_map::build_graph;
use ripvec_core::rerank::{DEFAULT_RERANK_CANDIDATES, DEFAULT_RERANK_MODEL, Reranker};
use serde::Deserialize;

const TOP_K: usize = 10;
const PAGERANK_ALPHA: f32 = 0.5;

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum RawTarget {
    Path(String),
    Span {
        path: String,
        #[serde(default)]
        start_line: Option<usize>,
        #[serde(default)]
        end_line: Option<usize>,
    },
}

#[derive(Debug, Clone)]
struct Target {
    path: String,
    start_line: Option<usize>,
    end_line: Option<usize>,
}

impl From<RawTarget> for Target {
    fn from(raw: RawTarget) -> Self {
        match raw {
            RawTarget::Path(path) => Self {
                path,
                start_line: None,
                end_line: None,
            },
            RawTarget::Span {
                path,
                start_line,
                end_line,
            } => Self {
                path,
                start_line,
                end_line,
            },
        }
    }
}

#[derive(Debug, Deserialize)]
struct RawTask {
    query: String,
    #[serde(default)]
    relevant: Vec<RawTarget>,
    #[serde(default)]
    secondary: Vec<RawTarget>,
    #[serde(default)]
    category: Option<String>,
}

struct Task {
    query: String,
    targets: Vec<Target>,
    category: String,
}

fn infer_category(query: &str) -> &'static str {
    if !query.trim().contains(' ') {
        return "symbol";
    }
    let q = query.to_lowercase();
    if q.starts_with("how ") {
        "architecture"
    } else {
        "semantic"
    }
}

fn path_matches(file_path: &str, target_path: &str) -> bool {
    let f = file_path.replace('\\', "/");
    let t = target_path.replace('\\', "/");
    f == t || f.ends_with(&format!("/{t}")) || t.ends_with(&format!("/{f}"))
}

fn target_matches_chunk(chunk: &CodeChunk, target: &Target) -> bool {
    if !path_matches(&chunk.file_path, &target.path) {
        return false;
    }
    match (target.start_line, target.end_line) {
        (Some(ts), Some(te)) => !(chunk.end_line < ts || chunk.start_line > te),
        _ => true,
    }
}

fn dcg(rels: &[u8]) -> f64 {
    rels.iter()
        .enumerate()
        .map(|(i, &r)| f64::from(r) / ((i + 2) as f64).log2())
        .sum()
}

fn ndcg_at_k(ranks: &[usize], n_relevant: usize, k: usize) -> f64 {
    if n_relevant == 0 {
        return 0.0;
    }
    let mut rels = vec![0u8; k];
    for &r in ranks {
        if (1..=k).contains(&r) {
            rels[r - 1] = 1;
        }
    }
    let ideal = dcg(&vec![1u8; k.min(n_relevant)]);
    if ideal > 0.0 { dcg(&rels) / ideal } else { 0.0 }
}

fn percentile(sorted: &[f64], p: f64) -> f64 {
    if sorted.is_empty() {
        return 0.0;
    }
    let n = sorted.len();
    let pos = (p / 100.0) * ((n - 1) as f64);
    let lo = pos.floor() as usize;
    let hi = pos.ceil() as usize;
    if lo == hi {
        return sorted[lo];
    }
    let frac = pos - lo as f64;
    sorted[lo] * (1.0 - frac) + sorted[hi] * frac
}

fn main() -> anyhow::Result<()> {
    let mut args = std::env::args().skip(1);
    let corpus: PathBuf = args
        .next()
        .ok_or_else(|| {
            anyhow::anyhow!("usage: corpus_bench <corpus_root> <annotations.json> [opts]")
        })?
        .into();
    let ann_path: PathBuf = args
        .next()
        .ok_or_else(|| anyhow::anyhow!("missing annotations.json"))?
        .into();

    let mut model_repo = DEFAULT_MODEL_REPO.to_string();
    let mut rerank_model = DEFAULT_RERANK_MODEL.to_string();
    let mut scope = Scope::All;
    let mut rerank_override: Option<bool> = None;
    let mut repeats: usize = 5;
    let mut candidate_k: usize = DEFAULT_RERANK_CANDIDATES;
    let mut args: Vec<String> = args.collect();
    let mut i = 0;
    while i < args.len() {
        match args[i].as_str() {
            "--model" => {
                args.remove(i);
                if i < args.len() {
                    model_repo = args.remove(i);
                }
            }
            "--rerank-model" => {
                args.remove(i);
                if i < args.len() {
                    rerank_model = args.remove(i);
                }
            }
            "--scope" => {
                args.remove(i);
                if i < args.len() {
                    scope = match args.remove(i).as_str() {
                        "code" => Scope::Code,
                        "docs" => Scope::Docs,
                        "all" => Scope::All,
                        o => anyhow::bail!("bad scope: {o}"),
                    };
                }
            }
            "--no-rerank" => {
                rerank_override = Some(false);
                args.remove(i);
            }
            "--rerank" => {
                rerank_override = Some(true);
                args.remove(i);
            }
            "--repeats" => {
                args.remove(i);
                if i < args.len() {
                    repeats = args.remove(i).parse()?;
                }
            }
            "--candidates" => {
                args.remove(i);
                if i < args.len() {
                    candidate_k = args.remove(i).parse()?;
                }
            }
            _ => i += 1,
        }
    }

    let raw: Vec<RawTask> = serde_json::from_slice(&std::fs::read(&ann_path)?)?;
    let tasks: Vec<Task> = raw
        .into_iter()
        .map(|t| Task {
            category: t
                .category
                .unwrap_or_else(|| infer_category(&t.query).to_string()),
            query: t.query,
            targets: t
                .relevant
                .into_iter()
                .chain(t.secondary)
                .map(Into::into)
                .collect(),
        })
        .collect();

    eprintln!("loading encoder ({model_repo}) + reranker ({rerank_model})...");
    let t0 = Instant::now();
    let encoder = StaticEncoder::from_pretrained(&model_repo)?;
    let reranker = Arc::new(Reranker::from_pretrained(&rerank_model)?);
    eprintln!("  loaded in {} ms", t0.elapsed().as_millis());

    // Build the repo graph + PageRank lookup. This is the MCP runtime
    // path *should* enable but currently doesn't (build_hybrid_from_ripvec
    // passes None); we're benchmarking against it being on so we measure
    // the real end-to-end power.
    eprintln!("building repo graph for PageRank...");
    let t0 = Instant::now();
    let pr_result = build_graph(&corpus);
    let (pr_lookup, pr_enabled) = match pr_result {
        Ok(graph) => {
            let lookup = pagerank_lookup(&graph);
            eprintln!(
                "  graph: {} files, pagerank lookup: {} entries",
                graph.files.len(),
                lookup.len()
            );
            let enabled = !lookup.is_empty();
            (lookup, enabled)
        }
        Err(e) => {
            eprintln!("  repo graph failed ({e}); continuing without PageRank");
            (std::collections::HashMap::new(), false)
        }
    };
    eprintln!("  pagerank build: {} ms", t0.elapsed().as_millis());

    eprintln!("building RipvecIndex for {}...", corpus.display());
    let t0 = Instant::now();
    let cfg = SearchConfig {
        scope,
        ..SearchConfig::default()
    };
    let profiler = Profiler::noop();
    let index = RipvecIndex::from_root(
        &corpus,
        encoder,
        &cfg,
        &profiler,
        if pr_enabled { Some(pr_lookup) } else { None },
        if pr_enabled { PAGERANK_ALPHA } else { 0.0 },
    )?;
    let index_ms = t0.elapsed().as_secs_f64() * 1000.0;
    let corpus_class = index.corpus_class();
    eprintln!(
        "  built in {index_ms:.0} ms ({} chunks, corpus={:?}, pr={})",
        index.chunks().len(),
        corpus_class,
        pr_enabled
    );

    // Per-query measurement loop. Each query runs `repeats` times to
    // capture warm-cache latency; we report median + p50/p90/p99 over
    // the per-query medians.
    let mut median_latencies: Vec<f64> = Vec::with_capacity(tasks.len());
    let mut q_ndcg10s: Vec<f64> = Vec::with_capacity(tasks.len());
    let mut q_recall10s: Vec<f64> = Vec::with_capacity(tasks.len());
    let mut q_precision10s: Vec<f64> = Vec::with_capacity(tasks.len());
    let mut by_cat: std::collections::BTreeMap<String, Vec<f64>> =
        std::collections::BTreeMap::new();
    let mut rerank_fired = 0usize;

    for task in &tasks {
        // Same gate as MCP run_search_ripvec uses.
        let auto_rerank = !is_symbol_query(&task.query)
            && match scope {
                Scope::Code => false,
                Scope::Docs => true,
                Scope::All => corpus_class.rerank_eligible(),
            };
        let do_rerank = rerank_override.map_or(auto_rerank, |f| f && !is_symbol_query(&task.query));
        if do_rerank {
            rerank_fired += 1;
        }

        let mut latencies = Vec::with_capacity(repeats);
        let mut ranked: Vec<(usize, f32)> = Vec::new();
        for _ in 0..repeats {
            let t = Instant::now();
            ranked = index.search(
                &task.query,
                candidate_k,
                SearchMode::Hybrid,
                None,
                None,
                None,
            );
            if do_rerank {
                let layer =
                    CrossEncoderRerank::new(reranker.clone(), task.query.clone(), candidate_k);
                let layers: Vec<Box<dyn RankingLayer>> = vec![Box::new(layer)];
                apply_chain(&mut ranked, index.chunks(), &layers);
            }
            latencies.push(t.elapsed().as_secs_f64() * 1000.0);
        }
        latencies.sort_by(|a, b| a.partial_cmp(b).unwrap());
        median_latencies.push(latencies[latencies.len() / 2]);

        let top_k: Vec<&CodeChunk> = ranked
            .iter()
            .take(TOP_K)
            .filter_map(|(idx, _)| index.chunks().get(*idx))
            .collect();
        let n_rel = task.targets.len();
        let mut ranks: Vec<usize> = Vec::new();
        for target in &task.targets {
            for (i, c) in top_k.iter().enumerate() {
                if target_matches_chunk(c, target) {
                    ranks.push(i + 1);
                    break;
                }
            }
        }
        let hits = ranks.len();
        let q_ndcg = ndcg_at_k(&ranks, n_rel, TOP_K);
        let q_recall = if n_rel == 0 {
            0.0
        } else {
            hits as f64 / n_rel as f64
        };
        let q_precision = hits as f64 / TOP_K as f64;
        q_ndcg10s.push(q_ndcg);
        q_recall10s.push(q_recall);
        q_precision10s.push(q_precision);
        by_cat
            .entry(task.category.clone())
            .or_default()
            .push(q_ndcg);
    }

    median_latencies.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let p50 = percentile(&median_latencies, 50.0);
    let p90 = percentile(&median_latencies, 90.0);
    let p95 = percentile(&median_latencies, 95.0);
    let p99 = percentile(&median_latencies, 99.0);
    let n = tasks.len() as f64;
    let mean = |xs: &[f64]| xs.iter().sum::<f64>() / n;

    let by_cat_summary: serde_json::Value = by_cat
        .iter()
        .map(|(k, v)| {
            (
                k.clone(),
                serde_json::json!({"n": v.len(), "ndcg10": v.iter().sum::<f64>() / v.len() as f64}),
            )
        })
        .collect::<serde_json::Map<_, _>>()
        .into();

    let out = serde_json::json!({
        "corpus": corpus.display().to_string(),
        "annotations": ann_path.display().to_string(),
        "model": model_repo,
        "scope": format!("{scope:?}").to_lowercase(),
        "corpus_class": format!("{corpus_class:?}").to_lowercase(),
        "tasks": tasks.len(),
        "rerank_fired": rerank_fired,
        "chunks": index.chunks().len(),
        "pagerank_enabled": pr_enabled,
        "index_ms": index_ms,
        "ndcg10": mean(&q_ndcg10s),
        "recall10": mean(&q_recall10s),
        "precision10": mean(&q_precision10s),
        "p50_ms": p50, "p90_ms": p90, "p95_ms": p95, "p99_ms": p99,
        "by_category": by_cat_summary,
    });
    println!("{}", serde_json::to_string_pretty(&out)?);
    Ok(())
}
