mod cli;
mod output;
mod progress;

use anyhow::{Context, Result};
use clap::Parser;
use tracing_subscriber::layer::SubscriberExt;

// Use mimalloc as the global allocator.
//
// Why: ripvec's hot paths — chunking, BM25 token-frequency maps, embedding
// `Vec<f32>` collection, search-time scoring — all churn many small/medium
// allocations across rayon worker threads. The default macOS libmalloc
// (libsystem_malloc / xzm_*) showed up as ~6% self in CPU profiles
// dominated by `_xzm_free`; mimalloc's per-thread heap-arena design is a
// significantly better fit for this allocation pattern (multi-threaded
// short-lived allocations with high locality) and reduces both wall-time
// allocator overhead and cross-thread contention.
//
// Set on the binary crate (not ripvec-core) because `#[global_allocator]`
// is a per-binary concern — downstream consumers of ripvec-core get to
// pick their own allocator.
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

/// Map `-v` count to a tracing level.
fn log_level_from_verbosity(verbose: u8) -> Option<&'static str> {
    match verbose {
        0 => None,
        1 => Some("info"),
        2 => Some("debug"),
        _ => Some("trace"),
    }
}

/// Resolve the effective tracing level.
///
/// Priority: explicit `--log-level`, `--debug`, `-v`, `RIPVEC_LOG`, then warn.
fn resolve_log_level(args: &cli::Args) -> String {
    args.log_level
        .clone()
        .or_else(|| args.debug.then(|| "debug".to_string()))
        .or_else(|| log_level_from_verbosity(args.verbose).map(str::to_string))
        .or_else(|| std::env::var("RIPVEC_LOG").ok())
        .unwrap_or_else(|| "warn".to_string())
}

/// Whether stderr should be line-oriented diagnostics instead of a spinner.
fn diagnostics_mode(args: &cli::Args) -> bool {
    args.profile || args.debug || args.verbose > 0
}

/// Build the ranking layer chain for a search.
///
/// Composes: PageRank boost (when a `RepoGraph` is available) ->
/// cross-encoder rerank. Rerank is always applied — the 4-corpus
/// benchmark matrix showed it lifts quality without measurable
/// per-query latency cost on warm cache. If the reranker model fails
/// to load (e.g., no network on first run), the layer is skipped with
/// a stderr warning and the bi-encoder ordering stands.
///
/// The reranker is loaded lazily and cached in `state` so subsequent
/// queries in the same process (e.g., TUI mode) reuse the model.
/// `state` is `&mut Option<Arc<Reranker>>` to allow mutation across
/// calls without re-cloning the args.
///
/// Returns `(layers, used_rerank)` so the caller can log timing
/// information if the rerank fired.
fn build_ranking_layers(
    args: &cli::Args,
    query: &str,
    pagerank: Option<(std::collections::HashMap<String, f32>, f32)>,
    reranker_cache: &mut Option<std::sync::Arc<ripvec_core::rerank::Reranker>>,
) -> (Vec<Box<dyn ripvec_core::ranking::RankingLayer>>, bool) {
    let mut layers: Vec<Box<dyn ripvec_core::ranking::RankingLayer>> = Vec::new();

    if let Some((pr, alpha)) = pagerank {
        layers.push(Box::new(ripvec_core::ranking::PageRankBoost::new(
            std::sync::Arc::new(pr),
            alpha,
        )));
    }

    // Lazy-load the reranker on first use; reuse across queries in
    // the same process.
    let reranker = if let Some(r) = reranker_cache.as_ref() {
        r.clone()
    } else {
        let repo = args
            .rerank_model
            .clone()
            .unwrap_or_else(|| ripvec_core::rerank::DEFAULT_RERANK_MODEL.to_string());
        match ripvec_core::rerank::Reranker::from_pretrained(&repo) {
            Ok(r) => {
                let arc = std::sync::Arc::new(r);
                *reranker_cache = Some(arc.clone());
                arc
            }
            Err(e) => {
                eprintln!(
                    "warning: reranker '{repo}' failed to load ({e}); \
                     continuing without rerank"
                );
                return (layers, false);
            }
        }
    };
    layers.push(Box::new(ripvec_core::ranking::CrossEncoderRerank::new(
        reranker,
        query.to_string(),
        args.rerank_candidates,
    )));
    (layers, true)
}

fn main() -> Result<()> {
    let args = cli::Args::parse();

    // Resolve model repo: the ripvec static encoder is the sole engine.
    let model_repo = cli::DEFAULT_MODEL_REPO;

    // Set up tracing: stderr (or --log-file), optional Chrome trace, level filter.
    let level = resolve_log_level(&args);
    let make_filter = || {
        tracing_subscriber::EnvFilter::try_new(&level)
            .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn"))
    };

    let trace_guard = args.trace.as_ref().map(|trace_path| {
        let (chrome_layer, guard) = tracing_chrome::ChromeLayerBuilder::new()
            .file(trace_path)
            .include_args(true)
            .build();
        tracing::subscriber::set_global_default(
            tracing_subscriber::registry()
                .with(make_filter())
                .with(chrome_layer),
        )
        .expect("failed to set tracing subscriber");
        guard
    });

    if trace_guard.is_none() {
        if let Some(log_path) = args.log_file.as_ref() {
            let fmt_layer = tracing_subscriber::fmt::layer()
                .with_target(false)
                .compact();
            let file = std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(log_path)
                .expect("failed to open log file");
            let _ = tracing::subscriber::set_global_default(
                tracing_subscriber::registry()
                    .with(make_filter())
                    .with(fmt_layer.with_writer(std::sync::Mutex::new(file))),
            );
        } else {
            let fmt_layer = tracing_subscriber::fmt::layer()
                .with_target(false)
                .compact();
            let _ = tracing::subscriber::set_global_default(
                tracing_subscriber::registry()
                    .with(make_filter())
                    .with(fmt_layer.with_writer(std::io::stderr)),
            );
        }
    }

    tracing::info!(
        path = %args.path,
        model = %model_repo,
        mode = %args.mode,
        threads = args.threads,
        verbose = args.verbose,
        "ripvec starting"
    );

    // Create profiler. --debug and -v modes enable phase diagnostics because
    // tracing alone does not show which long-running pipeline step is active.
    let profiler = ripvec_core::profile::Profiler::new(
        diagnostics_mode(&args),
        std::time::Duration::from_secs_f64(args.profile_interval),
    );

    // Configure thread pool — default to physical core count (empirically optimal)
    let cores = std::thread::available_parallelism()
        .map(std::num::NonZero::get)
        .unwrap_or(1);
    let threads = if args.threads > 0 {
        args.threads
    } else {
        cores
    };
    rayon::ThreadPoolBuilder::new()
        .num_threads(threads)
        .build_global()
        .context("failed to configure thread pool")?;

    // Whether to show indicatif progress bars. Diagnostic modes deliberately
    // use line-oriented output because spinners and logs share stderr.
    let use_progress =
        !diagnostics_mode(&args) && std::io::IsTerminal::is_terminal(&std::io::stderr());

    // Print profiler header
    profiler.header(
        env!("CARGO_PKG_VERSION"),
        model_repo,
        rayon::current_num_threads(),
        cores,
    );

    let search_cfg = build_search_config(&args);
    let result = run_oneshot_ripvec(&args, &search_cfg, use_progress, &profiler, model_repo);
    drop(trace_guard);
    result
}

/// Build a [`SearchConfig`] from CLI args without loading any model.
///
/// Used by the ripvec engine path to keep the walk/chunk parameters in
/// one place. The batch_size and max_tokens fields are not used by the
/// ripvec static encoder and are left at their defaults.
fn build_search_config(args: &cli::Args) -> ripvec_core::embed::SearchConfig {
    let mode: ripvec_core::hybrid::SearchMode = args.mode.parse().unwrap_or_default();
    build_search_config_with_mode(args, mode)
}

fn build_search_config_with_mode(
    args: &cli::Args,
    mode: ripvec_core::hybrid::SearchMode,
) -> ripvec_core::embed::SearchConfig {
    let mut search_cfg = ripvec_core::embed::SearchConfig {
        // batch_size and max_tokens are transformer-era fields; the ripvec
        // static encoder does not use them. Leave at defaults.
        batch_size: 0,
        max_tokens: 0,
        chunk: ripvec_core::chunk::ChunkConfig {
            max_chunk_bytes: args.max_chunk_bytes,
            window_size: args.window_size,
            window_overlap: args.window_overlap,
        },
        text_mode: args.text_mode,
        cascade_dim: None,
        file_type: args.file_type.clone(),
        exclude_extensions: args.exclude_extensions.clone(),
        include_extensions: args.include_extensions.clone(),
        ignore_patterns: Vec::new(),
        corpus: args.scope.into(),
        mode,
    };
    search_cfg.apply_repo_config(std::path::Path::new(&args.path));
    search_cfg
}

/// One-shot search through the ripvec static encoder path.
///
/// Loads `StaticEncoder` via `model2vec-rs`, builds a `RipvecIndex`
/// (walk + chunk via the ripvec engine's AST-merge chunker + encode + BM25
/// index), dispatches search by mode, and prints results via the
/// shared output formatter.
///
/// The ripvec engine's static embeddings rebuild on each invocation since
/// they are fast enough not to benefit from disk caching.
fn run_oneshot_ripvec(
    args: &cli::Args,
    search_cfg: &ripvec_core::embed::SearchConfig,
    use_progress: bool,
    profiler: &ripvec_core::profile::Profiler,
    model_repo: &str,
) -> Result<()> {
    anyhow::ensure!(!args.query.is_empty(), "query is required");

    let pb = use_progress.then(|| progress::spinner("Loading semble model\u{2026}"));
    let encoder = {
        let _guard = profiler.phase("model_load");
        ripvec_core::encoder::ripvec::dense::StaticEncoder::from_pretrained(model_repo)
            .context("failed to load model2vec model")?
    };
    if let Some(pb) = &pb {
        pb.finish_and_clear();
    }

    let pb = use_progress.then(|| progress::spinner("Indexing repo\u{2026}"));
    let index = ripvec_core::encoder::ripvec::index::RipvecIndex::from_root(
        std::path::Path::new(&args.path),
        encoder,
        search_cfg,
        profiler,
        // PageRank lookup is opt-in for the semble path. Building the
        // RepoGraph would re-walk the corpus through ripvec's
        // tree-sitter chunker — kept opt-out for now to preserve
        // semble's "all on CPU, sub-second" headroom. Layering it in
        // is a one-call follow-up once a sensible default is decided.
        None,
        0.0,
    )?;
    if let Some(pb) = &pb {
        pb.finish_and_clear();
    }

    let mode: ripvec_core::hybrid::SearchMode = args.mode.parse().unwrap_or_default();
    let corpus_size = index.chunks().len().max(1);
    let effective_top_k = if args.top_k > 0 {
        args.top_k.min(corpus_size)
    } else {
        corpus_size
    };
    tracing::info!(
        corpus_size,
        requested_top_k = args.top_k,
        effective_top_k,
        mode = ?mode,
        "ranking semble search results"
    );

    let mut ranked = index.search(
        &args.query,
        effective_top_k,
        mode,
        None, // auto-detect alpha via resolve_alpha
        None, // language filtering deferred (mapping is by extension)
        None, // path filtering deferred
    );

    // Apply the shared ranking-layer chain. The semble path doesn't
    // build a RepoGraph, so PageRank is None here; the only layer
    // that fires is cross-encoder rerank (when `--rerank` opts in
    // or auto-detect says yes). This is the path that closes the
    // bi-encoder-quality gap for prose corpora.
    if !ranked.is_empty() {
        let mut reranker_cache: Option<std::sync::Arc<ripvec_core::rerank::Reranker>> = None;
        let (layers, _) = build_ranking_layers(args, &args.query, None, &mut reranker_cache);
        if !layers.is_empty() {
            ripvec_core::ranking::apply_chain(&mut ranked, index.chunks(), &layers);
        }
    }

    let results: Vec<ripvec_core::embed::SearchResult> = ranked
        .into_iter()
        .filter_map(|(idx, score)| {
            index
                .chunks()
                .get(idx)
                .map(|chunk| ripvec_core::embed::SearchResult {
                    chunk: chunk.clone(),
                    similarity: score,
                })
        })
        .collect();

    profiler.finish();
    output::print_results(&results, &args.format);
    Ok(())
}
