//! `ckg search` — BM25 (default) / semantic / substring search over symbols.
use std::collections::{BTreeMap, HashMap};
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::Result;
use clap::Args as ClapArgs;
use cozo::DataValue;

#[derive(ClapArgs, Debug)]
pub struct Args {
    pub query: String,
    /// Use semantic (vector cosine) search; requires `Embedding` rows
    /// populated by a previous `ckg index` run (no `--no-embed`).
    #[arg(long)]
    pub semantic: bool,
    /// Use plain substring scan instead of BM25 ranking. Faster on huge
    /// codebases; less precise.
    #[arg(long)]
    pub substring: bool,
    #[arg(long, default_value_t = 20)]
    pub limit: u32,
    /// Include `@external` symbols (third-party / vendored code that
    /// SCIP persisted via `scip_verify::persist_external_symbols`) in
    /// results. Off by default — users want THEIR code first.
    #[arg(long)]
    pub include_external: bool,
}

pub fn run(home: Option<PathBuf>, args: Args) -> Result<()> {
    let (_base, storage) = super::repo_lookup::open_for(home, None)?;

    // #15: clamp `--limit 0` to 1. `as usize` would otherwise pass 0
    // through to `:limit 0` (empty Datalog result) or `truncate(0)`.
    let limit = (args.limit.max(1)) as usize;
    if args.semantic {
        return semantic_search(&storage, &args.query, limit, args.include_external);
    }
    if args.substring {
        return substring_search(&storage, &args.query, limit, args.include_external);
    }
    bm25_search(&storage, &args.query, limit, args.include_external)
}

// ----- BM25 ----------------------------------------------------------------

/// Standard Okapi BM25 over `name + qname + doc`. Loads all symbols into
/// memory once; fine ≤50k symbols. Scores: IDF(t) × TF(t,d)·(k1+1) /
/// (TF(t,d) + k1·(1 − b + b·|d|/avgdl)).
fn bm25_search(
    storage: &ckg_storage::Storage,
    query: &str,
    limit: usize,
    include_external: bool,
) -> Result<()> {
    const K1: f32 = 1.5;
    const B: f32 = 0.75;

    let q_terms = tokenize(query);
    if q_terms.is_empty() {
        println!("(empty query)");
        return Ok(());
    }

    // #14: BM25 loads every symbol row into memory; cap at MAX_BM25_DOCS
    // to bound process memory on huge monorepos. The fn comment above
    // calls 50k "fine" — at 200k we're at ~50 MB allocated for the
    // Doc vec, which is the upper bound we accept. Above the cap the
    // user should switch to `--semantic` (HNSW) or `--substring` (no
    // scoring, server-side filter).
    const MAX_BM25_DOCS: usize = 200_000;
    let rows = storage.run_immutable(&format!(
        "?[id, qname, name, kind, file, line, doc] := \
         *Symbol{{id, qname, name, kind, file, line, doc}} :limit {MAX_BM25_DOCS}"
    ))?;
    if rows.rows.len() >= MAX_BM25_DOCS {
        eprintln!(
            "(warning: BM25 search capped at {MAX_BM25_DOCS} symbols; \
             use `--semantic` or narrow with substring filters)"
        );
    }
    if rows.rows.is_empty() {
        println!("(no symbols indexed)");
        return Ok(());
    }

    // CLI-H10: intern terms via an `Arc<str>` pool. At 200k docs × ~30
    // terms each, the ~6M token strings contain massive duplication
    // (e.g. 50k docs containing the term "id"). Interning reduces
    // peak RSS from ~70MB to ~25MB on a representative monorepo.
    let mut pool: HashMap<String, Arc<str>> = HashMap::new();
    let intern = |pool: &mut HashMap<String, Arc<str>>, s: String| -> Arc<str> {
        if let Some(existing) = pool.get(&s) {
            existing.clone()
        } else {
            let arc: Arc<str> = Arc::from(s.as_str());
            pool.insert(s, arc.clone());
            arc
        }
    };

    struct Doc {
        id: String,
        qname: String,
        kind: String,
        file: String,
        line: String,
        terms: HashMap<Arc<str>, u32>,
        len: u32,
    }
    let mut docs: Vec<Doc> = Vec::with_capacity(rows.rows.len());
    let s = |v: &DataValue| match v {
        DataValue::Str(s) => s.to_string(),
        DataValue::Num(n) => format!("{n:?}"),
        _ => String::new(),
    };
    let mut total_len = 0u64;
    for r in &rows.rows {
        if r.len() < 7 {
            continue;
        }
        let qname = s(&r[1]);
        let name = s(&r[2]);
        let doc = s(&r[6]);
        // @external filter: skip vendored / third-party SCIP-persisted
        // symbols unless explicitly opted in.
        if !include_external && doc.contains("@external") {
            continue;
        }
        let combined = format!("{name} {qname} {doc}");
        let tokens = tokenize(&combined);
        let mut tf: HashMap<Arc<str>, u32> = HashMap::new();
        for t in tokens {
            let interned = intern(&mut pool, t);
            *tf.entry(interned).or_default() += 1;
        }
        let len = tf.values().sum::<u32>();
        total_len += len as u64;
        docs.push(Doc {
            id: s(&r[0]),
            qname,
            kind: s(&r[3]),
            file: s(&r[4]),
            line: s(&r[5]),
            terms: tf,
            len,
        });
    }
    let n = docs.len() as f32;
    let avgdl = if docs.is_empty() {
        1.0
    } else {
        (total_len as f32) / n
    };

    // IDF per query term (using "BM25 plus" smoothing).
    let mut idf: HashMap<Arc<str>, f32> = HashMap::new();
    for q in &q_terms {
        let q_arc = intern(&mut pool, q.clone());
        let df = docs.iter().filter(|d| d.terms.contains_key(&q_arc)).count() as f32;
        // ln((N − df + 0.5) / (df + 0.5) + 1) — always positive.
        idf.insert(q_arc, ((n - df + 0.5) / (df + 0.5) + 1.0).ln());
    }

    // Score every doc; collect non-zero.
    let mut scored: Vec<(f32, &Doc)> = Vec::new();
    for d in &docs {
        let mut score = 0.0f32;
        for q in &q_terms {
            let q_arc = pool.get(q.as_str());
            let tf = q_arc
                .and_then(|qa| d.terms.get(qa))
                .copied()
                .unwrap_or(0) as f32;
            if tf == 0.0 {
                continue;
            }
            let idf_q = q_arc
                .and_then(|qa| idf.get(qa))
                .copied()
                .unwrap_or(0.0);
            let dl_norm = 1.0 - B + B * (d.len as f32) / avgdl;
            let weight = idf_q * (tf * (K1 + 1.0)) / (tf + K1 * dl_norm);
            score += weight;
        }
        if score > 0.0 {
            scored.push((score, d));
        }
    }
    scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
    scored.truncate(limit);

    if scored.is_empty() {
        println!("(no matches)");
        return Ok(());
    }
    println!("{:<8} {:<60} FILE:LINE", "SCORE", "QNAME");
    for (score, d) in &scored {
        let _ = (&d.kind, &d.id);
        println!("{:<8.3} {:<60} {}:{}", score, d.qname, d.file, d.line);
    }
    Ok(())
}

/// Lowercase + alphanumeric token split. Splits on `::`, `.`, `_`, whitespace,
/// and any non-alphanumeric punctuation.
fn tokenize(s: &str) -> Vec<String> {
    s.split(|c: char| !c.is_alphanumeric())
        .filter(|s| !s.is_empty())
        .map(|t| t.to_lowercase())
        .collect()
}

// ----- substring fallback --------------------------------------------------

fn substring_search(
    storage: &ckg_storage::Storage,
    query: &str,
    limit: usize,
    include_external: bool,
) -> Result<()> {
    let mut p = BTreeMap::new();
    p.insert("q".into(), DataValue::from(query));
    let script = if include_external {
        "?[qname, kind, file, line] := \
         *Symbol{qname, kind, file, line}, str_includes(qname, $q)"
    } else {
        "?[qname, kind, file, line] := \
         *Symbol{qname, kind, file, line, doc}, \
         str_includes(qname, $q), \
         not str_includes(doc, \"@external\")"
    };
    let rows = storage.run_with_immutable(script, p)?;
    for row in rows.rows.iter().take(limit) {
        println!("  {:?}", row[0]);
    }
    Ok(())
}

// ----- semantic ------------------------------------------------------------

fn semantic_search(
    storage: &ckg_storage::Storage,
    query: &str,
    limit: usize,
    include_external: bool,
) -> Result<()> {
    let embedder = ckg_embed::Embedder::new();
    let qvec = embedder
        .embed_one(query)
        .map_err(|e| anyhow::anyhow!("embed query: {e}"))?;

    // Prefer Cozo's HNSW index for sub-millisecond nearest-neighbor; on
    // older DBs (no index yet, e.g. `--no-embed` was used last run), fall
    // back to brute-force cosine in memory.
    let scored = match storage.hnsw_search(&qvec, limit) {
        Ok(hits) if !hits.is_empty() => hits,
        _ => {
            // Cap the brute-force fallback at 200k embeddings — past
            // that the load itself OOMs on 1M-symbol monorepos.
            const FALLBACK_CAP: usize = 200_000;
            let embeddings = storage.iter_embeddings_capped(FALLBACK_CAP)?;
            if embeddings.is_empty() {
                anyhow::bail!(
                    "no embeddings stored for this repo — rerun `ckg index .` (without --no-embed)"
                );
            }
            if embeddings.len() >= FALLBACK_CAP {
                eprintln!(
                    "(warning: brute-force semantic scan capped at {FALLBACK_CAP} rows; \
                     run `ckg index --rebuild-index` to enable HNSW for full coverage)"
                );
            }
            // H8: filter out dim-mismatch rows instead of burying them at
            // score 0.0. A None from cosine() means the stored embedding
            // has a different dimension than the query vector — silently
            // keeping them at 0.0 pollutes the ranking. Skip them and
            // emit one warn per search call with the aggregate count.
            let mut skipped: usize = 0;
            let mut s: Vec<(String, f32)> = embeddings
                .into_iter()
                .filter_map(|(id, v)| match ckg_embed::cosine(&qvec, &v) {
                    Some(score) => Some((id, score)),
                    None => {
                        skipped += 1;
                        None
                    }
                })
                .collect();
            if skipped > 0 {
                tracing::warn!(
                    skipped,
                    "semantic search: skipped {} embedding(s) with dimension \
                     mismatch (stored dim != query dim {}); re-run `ckg index` \
                     to rebuild embeddings",
                    skipped,
                    qvec.len()
                );
            }
            s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            s.truncate(limit);
            s
        }
    };
    // Make output shape compatible with the rest of this fn.
    let scored: Vec<(f32, String)> = scored.into_iter().map(|(id, s)| (s, id)).collect();

    if scored.is_empty() {
        println!("(no matches)");
        return Ok(());
    }
    let ids: Vec<DataValue> = scored
        .iter()
        .map(|(_, id)| DataValue::from(id.as_str()))
        .collect();
    let mut params = BTreeMap::new();
    params.insert("ids".into(), DataValue::List(ids));
    let rows = storage.run_with_immutable(
        "?[id, qname, kind, file, line, doc] := \
         *Symbol{id, qname, kind, file, line, doc}, is_in(id, $ids)",
        params,
    )?;
    let s = |v: &DataValue| match v {
        DataValue::Str(s) => s.to_string(),
        DataValue::Num(n) => format!("{n:?}"),
        _ => String::new(),
    };
    // (qname, kind, file, line, doc) — keep `doc` so we can filter
    // @external rows after the ranking pass.
    let info: HashMap<String, (String, String, String, String, String)> = rows
        .rows
        .into_iter()
        .filter_map(|r| {
            if r.len() < 6 {
                return None;
            }
            Some((
                s(&r[0]),
                (s(&r[1]), s(&r[2]), s(&r[3]), s(&r[4]), s(&r[5])),
            ))
        })
        .collect();

    println!("{:<8} {:<60} FILE:LINE", "SCORE", "QNAME");
    for (score, id) in &scored {
        let (q, _k, f, l, doc) = info
            .get(id)
            .cloned()
            .unwrap_or_else(|| {
                (
                    id.clone(),
                    String::new(),
                    String::new(),
                    String::new(),
                    String::new(),
                )
            });
        if !include_external && doc.contains("@external") {
            continue;
        }
        println!("{:<8.3} {:<60} {f}:{l}", score, q);
    }
    Ok(())
}
