//! Smart context assembly — given a question, return the minimal code set to answer it.
//!
//! Algorithm:
//! 1. Search for seed results
//! 2. BFS expand via call graph (callers/callees/both)
//! 3. Cap expansion at 200 nodes
//! 4. Deduplicate by parent_id
//! 5. Sort by file → line (reading order)

use std::collections::{HashMap, HashSet, VecDeque};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use rayon::prelude::*;

use crate::parser::{ChunkType, Language};
use crate::{AnalysisError, Embedder, Embedding};

use crate::store::helpers::{CallGraph, SearchFilter};
use crate::store::SearchResult;
use crate::Store;

/// Default maximum nodes in BFS expansion to prevent blowup on hub functions.
pub const DEFAULT_MAX_EXPANDED_NODES: usize = 200;

/// Options for gather operation
#[derive(Debug, Clone)]
pub struct GatherOptions {
    pub expand_depth: usize,
    pub direction: GatherDirection,
    pub limit: usize,
    pub seed_limit: usize,
    pub seed_threshold: f32,
    pub decay_factor: f32,
    /// Maximum nodes in BFS expansion (default: 200).
    /// Prevents blowup on hub functions with many callers/callees.
    pub max_expanded_nodes: usize,
    /// Pre-computed query embedding. When set, `gather()` uses this instead
    /// of calling `embedder.embed_query()`. Useful when the caller shares one
    /// embedding across multiple phases (e.g., `task()` reuses it for scout + gather).
    pub query_embedding: Option<Embedding>,
}

impl GatherOptions {
    /// Sets the maximum depth for expanding nested structures in output.
    ///
    /// # Arguments
    ///
    /// * `depth` - The maximum nesting level to expand before truncating or summarizing nested content.
    ///
    /// # Returns
    ///
    /// Returns `self` to allow method chaining.
    pub fn with_expand_depth(mut self, depth: usize) -> Self {
        self.expand_depth = depth;
        self
    }
    /// Sets the direction for gathering operations.
    ///
    /// # Arguments
    ///
    /// * `direction` - The direction to use when gathering data
    ///
    /// # Returns
    ///
    /// Returns `self` to allow method chaining.
    pub fn with_direction(mut self, direction: GatherDirection) -> Self {
        self.direction = direction;
        self
    }
    /// Sets the maximum number of items to process.
    ///
    /// # Arguments
    ///
    /// * `limit` - The maximum number of items to process
    ///
    /// # Returns
    ///
    /// Returns `self` to allow method chaining.
    pub fn with_limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }
    /// Sets the maximum number of seeds to process.
    ///
    /// # Arguments
    ///
    /// * `limit` - The maximum number of seeds to process
    ///
    /// # Returns
    ///
    /// Returns `self` to allow method chaining.
    pub fn with_seed_limit(mut self, limit: usize) -> Self {
        self.seed_limit = limit;
        self
    }
    /// Sets the seed threshold for this builder and returns self for method chaining.
    ///
    /// # Arguments
    ///
    /// * `threshold` - The seed threshold value (typically between 0.0 and 1.0) used to filter seeds during processing
    ///
    /// # Returns
    ///
    /// Returns `Self` to allow for method chaining in builder pattern style construction.
    pub fn with_seed_threshold(mut self, threshold: f32) -> Self {
        // RT-RES-7: Guard against NaN — use default (0.3) if threshold is not finite.
        if threshold.is_finite() {
            self.seed_threshold = threshold;
        } else {
            tracing::warn!(
                threshold = %threshold,
                "NaN/infinite seed threshold, using default 0.3"
            );
            self.seed_threshold = 0.3;
        }
        self
    }
    /// Sets the decay factor for this builder, clamping the value to the range [0.0, 1.0].
    ///
    /// If the provided factor is not a finite number (NaN or infinite), the current decay factor is preserved unchanged.
    ///
    /// # Arguments
    ///
    /// * `factor` - The decay factor to set, typically a value between 0.0 and 1.0
    ///
    /// # Returns
    ///
    /// Returns `self` to allow for method chaining.
    pub fn with_decay_factor(mut self, factor: f32) -> Self {
        self.decay_factor = if factor.is_finite() {
            factor.clamp(0.0, 1.0)
        } else {
            self.decay_factor
        };
        self
    }
    /// Sets the maximum number of nodes that can be expanded during search.
    ///
    /// # Arguments
    ///
    /// * `max` - The maximum number of nodes to expand. When the search expands more nodes than this limit, it will stop exploring further branches.
    ///
    /// # Returns
    ///
    /// Returns `self` to allow method chaining.
    pub fn with_max_expanded_nodes(mut self, max: usize) -> Self {
        self.max_expanded_nodes = max;
        self
    }
}

/// Read CQS_GATHER_MAX_NODES once via OnceLock, not on every GatherOptions::default().
///
/// Public so CLI text-mode warnings can report the actual cap. Parse/warn/
/// default goes through the shared `crate::limits::parse_env_usize`; the
/// OnceLock cache keeps it a single env read per process.
pub fn gather_max_nodes() -> usize {
    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *CAP.get_or_init(|| {
        crate::limits::parse_env_usize("CQS_GATHER_MAX_NODES", DEFAULT_MAX_EXPANDED_NODES)
    })
}

impl Default for GatherOptions {
    /// Creates a new instance with default configuration values for graph traversal and expansion.
    ///
    /// The default settings prioritize balanced exploration with a depth of 1, bidirectional search, a limit of 10 results, and a decay factor of 0.8 to gradually reduce node relevance as distance increases.
    ///
    /// # Returns
    ///
    /// A new `Self` instance initialized with sensible defaults for graph expansion operations.
    fn default() -> Self {
        Self {
            expand_depth: 1,
            direction: GatherDirection::Both,
            limit: 10,
            seed_limit: 5,
            seed_threshold: 0.3,
            decay_factor: 0.8,
            max_expanded_nodes: gather_max_nodes(),
            query_embedding: None,
        }
    }
}

/// Direction of call graph expansion
#[derive(
    Debug,
    Clone,
    Copy,
    PartialEq,
    Eq,
    serde::Serialize,
    serde::Deserialize,
    clap::ValueEnum,
    schemars::JsonSchema,
)]
// Lowercase on the wire/schema to match the CLI spelling (`both`/`callers`/
// `callees` via clap `FromStr`); without it serde + schemars would expect the
// PascalCase variant names.
#[serde(rename_all = "lowercase")]
pub enum GatherDirection {
    Both,
    Callers,
    Callees,
}

impl std::str::FromStr for GatherDirection {
    type Err = String;
    /// Parses a string into a direction value.
    ///
    /// Converts a string representation into the corresponding direction variant. Valid inputs are "both", "callers", or "callees" (case-sensitive).
    ///
    /// # Arguments
    ///
    /// * `s` - A string slice to parse as a direction
    ///
    /// # Returns
    ///
    /// Returns `Ok(Self)` with the parsed direction variant if the input is valid, or `Err(String)` with a descriptive error message if the input is not recognized.
    fn from_str(s: &str) -> std::result::Result<Self, String> {
        match s {
            "both" => Ok(Self::Both),
            "callers" => Ok(Self::Callers),
            "callees" => Ok(Self::Callees),
            _ => Err(format!(
                "Invalid direction '{}'. Valid: both, callers, callees",
                s
            )),
        }
    }
}

/// A gathered code chunk with context
#[derive(Debug, Clone, serde::Serialize)]
pub struct GatheredChunk {
    pub name: String,
    #[serde(serialize_with = "crate::serialize_path_normalized")]
    pub file: PathBuf,
    pub line_start: u32,
    pub line_end: u32,
    pub language: Language,
    pub chunk_type: ChunkType,
    pub signature: String,
    pub content: String,
    pub score: f32,
    pub depth: usize,
    /// Source: None = project, Some(name) = reference
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source: Option<String>,
    /// Ranking provenance for seed chunks (skip-when-empty). Mirrors the
    /// `search` surface. Only seed hits (depth 0) carry it — BFS-expanded
    /// chunks were not retrieved by fusion, so they have no leg ranks or
    /// boosts and stay empty.
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub rank_signals: Vec<crate::store::RankSignal>,
}

impl GatheredChunk {
    /// Build from a `SearchResult`, relativizing its file path against `root`.
    pub fn from_search(
        sr: &crate::store::SearchResult,
        root: &Path,
        score: f32,
        depth: usize,
        source: Option<String>,
    ) -> Self {
        Self {
            name: sr.chunk.name.clone(),
            file: crate::relativize_or_warn(&sr.chunk.file, root),
            line_start: sr.chunk.line_start,
            line_end: sr.chunk.line_end,
            language: sr.chunk.language,
            chunk_type: sr.chunk.chunk_type,
            signature: sr.chunk.signature.clone(),
            content: sr.chunk.content.clone(),
            score,
            depth,
            source,
            rank_signals: Vec::new(),
        }
    }
}

/// Result of a gather operation
#[derive(Debug, Clone, serde::Serialize)]
pub struct GatherResult {
    pub chunks: Vec<GatheredChunk>,
    pub expansion_capped: bool,
    /// True if batch name search failed and results may be incomplete
    pub search_degraded: bool,
}

// MAX_EXPANDED_NODES is now configurable via GatherOptions::max_expanded_nodes
// (default: DEFAULT_MAX_EXPANDED_NODES = 200)

/// BFS-expand names via call graph, applying score decay and enforcing a node cap.
///
/// Returns `(name_scores, expansion_capped)` where `name_scores` maps
/// function names to `(score, depth)`.
pub(crate) fn bfs_expand(
    name_scores: &mut HashMap<String, (f32, usize)>,
    graph: &CallGraph,
    opts: &GatherOptions,
) -> bool {
    let mut expansion_capped = false;
    if opts.expand_depth == 0 {
        return false;
    }

    // Track visited nodes to prevent re-expansion from overlapping seeds.
    // Without this, if two seeds overlap (e.g., bridge results matching ref seeds
    // in gather_cross_index), the same node gets expanded twice, doubling neighbor lookups.
    // Use Arc<str> in visited/queue to avoid per-node String heap allocations.
    let mut visited: HashSet<Arc<str>> =
        name_scores.keys().map(|k| Arc::from(k.as_str())).collect();
    let initial_size = name_scores.len();

    // HashMap iteration order is process-seed-dependent, so enqueueing
    // straight from `name_scores.keys()` would make BFS non-deterministic:
    // when `max_expanded_nodes` fires mid-expansion, which seeds got their
    // neighbors in before the cap would depend on randomness. Sort by
    // (score desc, name asc) so a given input produces the same output on
    // every run, regardless of process-wide HashMap randomization.
    let mut seeds: Vec<(&str, f32)> = name_scores
        .iter()
        .map(|(k, (s, _))| (k.as_str(), *s))
        .collect();
    // Use `total_cmp` to match the rest of the search/scoring pipeline
    // (BoundedScoreHeap, MMR, apply_parent_boost re-sort,
    // search_across_projects merge). `partial_cmp().unwrap_or(Equal)`
    // collapses NaN to "equal to everyone", which makes the secondary
    // `name asc` tiebreak inherit the input slice's NaN position (not
    // stable under sort). `total_cmp` gives a deterministic total order
    // over all f32 values including NaN.
    seeds.sort_by(|(a_name, a_score), (b_name, b_score)| {
        b_score.total_cmp(a_score).then_with(|| a_name.cmp(b_name))
    });

    let mut queue: VecDeque<(Arc<str>, usize)> = VecDeque::new();
    for (name, _) in seeds {
        queue.push_back((Arc::from(name), 0));
    }

    while let Some((name, depth)) = queue.pop_front() {
        if depth >= opts.expand_depth {
            continue;
        }
        if name_scores.len() >= opts.max_expanded_nodes && visited.len() > initial_size {
            expansion_capped = true;
            break;
        }

        let neighbors = get_neighbors(graph, &name, opts.direction);
        let base_score = name_scores
            .get(name.as_ref())
            .map(|(s, _)| *s)
            .unwrap_or(0.5);
        let new_score = base_score * opts.decay_factor;
        for neighbor in neighbors {
            if !visited.contains(&neighbor) {
                // Cap is on `name_scores.len()`, so it only matters for new
                // insertions. Already-visited bumps below don't grow the map
                // — let them run unconditionally so a node reached late via a
                // higher-score path still gets its score promoted.
                if name_scores.len() >= opts.max_expanded_nodes {
                    expansion_capped = true;
                    break;
                }
                visited.insert(Arc::clone(&neighbor));
                let key: String = neighbor.to_string();
                name_scores.insert(key, (new_score, depth + 1));
                queue.push_back((neighbor, depth + 1));
            } else if let Some(existing) = name_scores.get_mut(neighbor.as_ref()) {
                // Already visited — update score if higher, preserve minimum depth.
                // We intentionally do NOT re-enqueue the node. Re-enqueueing
                // would propagate the higher score to neighbors (Dijkstra-like), but
                // risks exponential blowup on dense graphs. The max_expanded_nodes cap
                // keeps memory bounded, and the score update here ensures the node itself
                // gets the best available score even without re-expansion.
                //
                // This branch must run even when the cap has been hit — the
                // bump doesn't grow `name_scores`.
                //
                // Depth-min update runs unconditionally (not score-gated) so a
                // low-score-shorter-path (e.g. seed B score 0.01 reaching D in
                // depth 1) can update D's depth even when D was recorded at
                // depth 2 from a high-score-deeper-path (seed A score 1.0 →
                // A→C→D). The depth field surfaces in `GatheredChunk.depth` and
                // on the JSON output. Depth-min is idempotent and cheap;
                // running it always ensures the reported depth is the shortest
                // path's.
                existing.1 = existing.1.min(depth + 1);
                if new_score > existing.0 {
                    existing.0 = new_score;
                }
            }
        }
        if expansion_capped {
            break;
        }
    }
    expansion_capped
}

/// Batch-fetch chunks for expanded names, deduplicate by id, assemble `GatheredChunk`s.
///
/// `overlay_seeds`, when `Some`, maps an overlay-origin SEED name to its overlay
/// `SearchResult` (built by `WorktreeOverlay::overlay_origin_seeds` /
/// `overlay_seed_chunks_for_names`). A depth-0 name present there is assembled
/// from the OVERLAY chunk and EXCLUDED from the parent by-name batch, so a
/// worktree-added seed surfaces (the parent store has no chunk for it) and a
/// worktree-modified seed shows its worktree content (not the parent's stale
/// copy). Expansion (depth > 0) always re-fetches from the parent store —
/// graph expansion stays parent-truth in Part A. `None` (the no-overlay path
/// and the cross-index / onboard callers) is byte-identical to the prior
/// parent-only behavior.
///
/// Returns `(chunks, search_degraded)`.
pub(crate) fn fetch_and_assemble<Mode>(
    store: &Store<Mode>,
    name_scores: &HashMap<String, (f32, usize)>,
    root: &Path,
    overlay_seeds: Option<&HashMap<String, crate::store::SearchResult>>,
) -> (Vec<GatheredChunk>, bool) {
    // A name is served from the overlay only at depth 0. Exclude those names
    // from the parent batch so we never spend an FTS probe re-fetching a chunk
    // we're about to override (and so a parent name-collision can't sneak in).
    let is_overlay_seed = |name: &str, depth: usize| -> bool {
        depth == 0 && overlay_seeds.is_some_and(|seeds| seeds.contains_key(name))
    };

    let all_names: Vec<&str> = name_scores
        .iter()
        .filter(|(name, (_, depth))| !is_overlay_seed(name, *depth))
        .map(|(name, _)| name.as_str())
        .collect();
    let (batch_results, search_degraded) = match store.search_by_names_batch(&all_names, 1) {
        Ok(r) => (r, false),
        Err(e) => {
            tracing::warn!(error = %e, "Batch name search failed, results may be incomplete");
            (HashMap::new(), true)
        }
    };

    let mut seen_ids: HashSet<String> = HashSet::new();
    let mut chunks: Vec<GatheredChunk> = Vec::new();

    for (name, (score, depth)) in name_scores {
        // Prefer the overlay seed for an overlay-origin depth-0 name: the
        // parent store either has no chunk for it (worktree-added) or a stale
        // one (worktree-modified). This is the gather/task analogue of scout
        // consuming the merged `SearchResult`s directly.
        if is_overlay_seed(name, *depth) {
            if let Some(r) = overlay_seeds.and_then(|seeds| seeds.get(name)) {
                if seen_ids.insert(r.chunk.id.clone()) {
                    chunks.push(GatheredChunk::from_search(r, root, *score, *depth, None));
                }
                continue;
            }
        }
        if let Some(results) = batch_results.get(name) {
            if let Some(r) = results.first() {
                // HashSet::insert returns false if the key was already
                // present — single hash probe instead of two.
                if !seen_ids.insert(r.chunk.id.clone()) {
                    continue;
                }
                chunks.push(GatheredChunk::from_search(r, root, *score, *depth, None));
            }
        }
    }

    tracing::debug!(chunk_count = chunks.len(), "Chunks assembled");
    (chunks, search_degraded)
}

/// Attach seed-search `rank_signals` to the matching depth-0 chunks.
///
/// `seed_signals` maps a seed chunk name to the provenance recorded by the
/// seed fusion search. Only depth-0 (seed) chunks are eligible — BFS-expanded
/// chunks were fetched by name lookup, not fusion, so they have no leg ranks
/// or boosts. Keying on `(name, depth == 0)` mirrors how the seed scores were
/// inserted into `name_scores`.
pub(crate) fn attach_seed_signals(
    chunks: &mut [GatheredChunk],
    seed_signals: &HashMap<String, Vec<crate::store::RankSignal>>,
) {
    if seed_signals.is_empty() {
        return;
    }
    for c in chunks.iter_mut() {
        if c.depth == 0 {
            if let Some(sigs) = seed_signals.get(&c.name) {
                c.rank_signals = sigs.clone();
            }
        }
    }
}

/// Sort chunks by score desc (file/line/name tiebreak), truncate to limit,
/// then re-sort to file/line reading order.
///
/// The tiebreak cascade (name → file → line_start) is full — two chunks can
/// collide on `name` when they're from different files (e.g., two `new`
/// functions), so we also compare `file` and `line_start` to keep the
/// truncate() boundary deterministic across process invocations.
pub(crate) fn sort_and_truncate(chunks: &mut Vec<GatheredChunk>, limit: usize) {
    chunks.sort_by(|a, b| {
        b.score
            .total_cmp(&a.score)
            .then(a.name.cmp(&b.name))
            .then(a.file.cmp(&b.file))
            .then(a.line_start.cmp(&b.line_start))
    });
    chunks.truncate(limit);
    chunks.sort_by(|a, b| {
        a.file
            .cmp(&b.file)
            .then(a.line_start.cmp(&b.line_start))
            .then(a.name.cmp(&b.name))
    });
}

/// Gather relevant code chunks for a query.
///
/// Embeds the query internally (or uses `opts.query_embedding` if pre-computed).
/// Loads the call graph internally. For pre-loaded graph, use [`gather_with_graph`].
/// Serves the parent index for the seed search; for worktree-overlaid seeds use
/// [`gather_with_overlay`].
pub fn gather<Mode>(
    store: &Store<Mode>,
    embedder: &Embedder,
    description: &str,
    opts: &GatherOptions,
    root: &Path,
) -> Result<GatherResult, AnalysisError> {
    gather_with_overlay(store, embedder, description, opts, root, None)
}

/// Like [`gather`] but accepts an optional worktree search overlay (
/// Part A). `Some` shadows the parent index for the SEED search; `None` is
/// byte-identical to [`gather`]. Only the daemon path (from an eligible
/// worktree) passes `Some` — see `dispatch_gather`. The BFS expansion stays on
/// parent-truth (the adapter emits a `seed-only` `_meta.overlay_graph` marker).
pub fn gather_with_overlay<Mode>(
    store: &Store<Mode>,
    embedder: &Embedder,
    description: &str,
    opts: &GatherOptions,
    root: &Path,
    overlay: Option<&crate::worktree_overlay::WorktreeOverlay>,
) -> Result<GatherResult, AnalysisError> {
    let query_embedding = match &opts.query_embedding {
        Some(emb) => emb.clone(),
        None => embedder.embed_query(description)?,
    };
    let graph = store.get_call_graph()?;
    gather_with_graph_overlay(
        store,
        &query_embedding,
        description,
        opts,
        root,
        &graph,
        overlay,
    )
}

/// Like [`gather`] but accepts a pre-loaded call graph.
///
/// Use when the caller already has the graph (e.g., batch mode or `task()`
/// which shares the graph across phases). Serves the parent index for the seed
/// search; for worktree-overlaid seeds use [`gather_with_graph_overlay`].
pub fn gather_with_graph<Mode>(
    store: &Store<Mode>,
    query_embedding: &crate::Embedding,
    query_text: &str,
    opts: &GatherOptions,
    root: &Path,
    graph: &CallGraph,
) -> Result<GatherResult, AnalysisError> {
    gather_with_graph_overlay(store, query_embedding, query_text, opts, root, graph, None)
}

/// Like [`gather_with_graph`] but accepts an optional worktree search overlay
/// (Part A). `Some` shadows the parent index for the SEED search so a
/// worktree-added/edited file surfaces as a gather seed; `None` is byte-
/// identical to [`gather_with_graph`]. The BFS / call-graph expansion stays on
/// parent-truth — only the seed is overlaid (Part B extends the graph).
#[allow(clippy::too_many_arguments)]
pub fn gather_with_graph_overlay<Mode>(
    store: &Store<Mode>,
    query_embedding: &crate::Embedding,
    query_text: &str,
    opts: &GatherOptions,
    root: &Path,
    graph: &CallGraph,
    overlay: Option<&crate::worktree_overlay::WorktreeOverlay>,
) -> Result<GatherResult, AnalysisError> {
    let _span = tracing::info_span!(
        "gather",
        query_len = query_text.len(),
        expand_depth = opts.expand_depth,
        limit = opts.limit
    )
    .entered();

    // 1. Seed with hybrid RRF search (not raw embedding-only)
    let filter = SearchFilter {
        query_text: query_text.to_string(),
        enable_rrf: false, // RRF off by default — pure cosine is faster + higher R@1 on expanded eval
        // Record provenance for the seed hits — attached to depth-0 chunks
        // below. Side channel; no scoring change.
        record_rank_signals: true,
        ..SearchFilter::default()
    };
    let seed_results = store.search_filtered(
        query_embedding,
        &filter,
        opts.seed_limit,
        opts.seed_threshold,
    )?;

    // Worktree overlay (Part A): shadow the parent index for the SEED
    // search before BFS expansion. Inactive (`None`) ⇒ no-op. The expansion
    // below walks the parent call graph — only the seeds reflect the worktree
    // delta (the `seed-only` `_meta.overlay_graph` marker).
    let seed_results = match overlay {
        Some(ov) => ov.merge_seed_results(
            seed_results,
            query_embedding,
            &filter,
            opts.seed_limit,
            opts.seed_threshold,
        ),
        None => seed_results,
    };
    tracing::debug!(seed_count = seed_results.len(), "Seed search complete");
    if seed_results.is_empty() {
        return Ok(GatherResult {
            chunks: Vec::new(),
            expansion_capped: false,
            search_degraded: false,
        });
    }

    // Seed names with their scores, and a name → seed `rank_signals` map so the
    // depth-0 chunks can carry their provenance after assembly.
    let mut name_scores: HashMap<String, (f32, usize)> = HashMap::new();
    let mut seed_signals: HashMap<String, Vec<crate::store::RankSignal>> = HashMap::new();
    for r in &seed_results {
        name_scores.insert(r.chunk.name.clone(), (r.score, 0));
        if !r.rank_signals.is_empty() {
            seed_signals
                .entry(r.chunk.name.clone())
                .or_insert_with(|| r.rank_signals.clone());
        }
    }

    // Overlay-origin seeds: assemble these depth-0 chunks from the OVERLAY
    // store, not the parent. The parent has no chunk for a worktree-added seed
    // (it would be silently dropped) and a stale one for a worktree-modified
    // seed — so without this, the `seed-only` `_meta.overlay_graph` marker
    // would lie for the headline case (new code added in the lane). Only the
    // seed surfaces overlaid; the BFS expansion below stays parent-truth.
    let overlay_seeds = overlay.map(|ov| ov.overlay_origin_seeds(&seed_results));

    // 2. BFS expand
    let expansion_capped = bfs_expand(&mut name_scores, graph, opts);
    tracing::info!(
        expanded = name_scores.len(),
        capped = expansion_capped,
        "BFS expansion complete"
    );

    // 3. Batch-fetch chunks, deduplicate
    let (mut chunks, search_degraded) =
        fetch_and_assemble(store, &name_scores, root, overlay_seeds.as_ref());

    // 3b. Attach seed provenance to depth-0 chunks. Expanded chunks (depth > 0)
    // weren't retrieved by fusion and keep an empty `rank_signals`.
    attach_seed_signals(&mut chunks, &seed_signals);

    // 4. Sort by score desc, truncate to limit, re-sort to reading order
    sort_and_truncate(&mut chunks, opts.limit);

    tracing::info!(final_chunks = chunks.len(), "Gather complete");

    Ok(GatherResult {
        chunks,
        expansion_capped,
        search_degraded,
    })
}

/// Cross-index gather: seed from a reference index, bridge into project code, BFS expand.
///
/// Accepts an optional HNSW index for O(log n) bridge searches instead of
/// brute-force scans per reference seed (pass `None` to scan).
///
/// Flow:
/// 1. Search reference index for seed chunks matching the query
/// 2. Retrieve seed chunk embeddings from the reference store
/// 3. For each seed embedding, search the project store for similar code (bridge)
/// 4. BFS expand project-side bridges via the project call graph
/// 5. Return both reference seeds (context) and expanded project chunks
pub fn gather_cross_index_with_index<Mode: Sync>(
    project_store: &Store<Mode>,
    ref_idx: &crate::reference::ReferenceIndex,
    query_embedding: &crate::Embedding,
    query_text: &str,
    opts: &GatherOptions,
    root: &Path,
    project_index: Option<&dyn crate::index::VectorIndex>,
) -> Result<GatherResult, AnalysisError> {
    let _span = tracing::info_span!(
        "gather_cross_index",
        ref_name = %ref_idx.name,
        query_len = query_text.len(),
        expand_depth = opts.expand_depth,
        limit = opts.limit,
    )
    .entered();

    // Model compatibility check: warn if project and reference use different embedding models
    if let (Ok(proj_model), Ok(ref_model)) = (
        project_store.get_metadata("model_name"),
        ref_idx.store.get_metadata("model_name"),
    ) {
        if proj_model != ref_model {
            tracing::warn!(
                project = %proj_model,
                reference = %ref_model,
                "Model mismatch between project and reference — results may be inaccurate"
            );
        }
    }

    // 1. Seed search against reference index (unweighted — user explicitly targets this ref)
    let filter = crate::store::helpers::SearchFilter {
        query_text: query_text.to_string(),
        enable_rrf: false, // RRF off by default — pure cosine is faster + higher R@1 on expanded eval
        ..SearchFilter::default()
    };
    let ref_seeds = crate::reference::search_reference(
        ref_idx,
        query_embedding,
        &filter,
        opts.seed_limit,
        opts.seed_threshold,
        false, // no weight for cross-index gather (user explicitly targets this ref)
    )?;
    tracing::debug!(
        ref_seed_count = ref_seeds.len(),
        "Reference seed search complete"
    );

    if ref_seeds.is_empty() {
        return Ok(GatherResult {
            chunks: Vec::new(),
            expansion_capped: false,
            search_degraded: false,
        });
    }

    // Collect ref seed chunk IDs for embedding retrieval
    let ref_seed_ids: Vec<&str> = ref_seeds.iter().map(|r| r.chunk.id.as_str()).collect();

    // 2. Get embeddings for ref seed chunks
    let ref_embeddings = match ref_idx.store.get_embeddings_by_ids(&ref_seed_ids) {
        Ok(e) => e,
        Err(e) => {
            tracing::warn!(error = %e, "Failed to get ref seed embeddings, falling back to query embedding only");
            HashMap::new()
        }
    };

    // Build ref seed output chunks (these go into the result as reference context).
    // Use Path::new("") as root since ref files don't need path relativization.
    let ref_chunks: Vec<GatheredChunk> = ref_seeds
        .iter()
        .map(|r| {
            GatheredChunk::from_search(r, Path::new(""), r.score, 0, Some(ref_idx.name.clone()))
        })
        .collect();

    // 3. Bridge: for each ref seed, search the project store with the seed's embedding
    //    to find semantically similar project code.
    //    If no embedding available for a seed, use the original query embedding.
    //    Parallelized with rayon — Store is Send+Sync (SqlitePool + Runtime + AtomicBool).
    let bridge_filter = SearchFilter {
        query_text: query_text.to_string(),
        enable_rrf: false, // RRF off by default — pure cosine is faster + higher R@1 on expanded eval
        ..SearchFilter::default()
    };

    let bridge_limit = 3; // Top 3 project matches per ref seed

    let _bridge_span = tracing::info_span!("bridge_search", seed_count = ref_seeds.len()).entered();

    let bridge_error_count = std::sync::atomic::AtomicUsize::new(0);
    let bridge_results: Vec<(f32, Vec<SearchResult>)> = ref_seeds
        .par_iter()
        .filter_map(|seed| {
            let search_embedding = ref_embeddings
                .get(&seed.chunk.id)
                .unwrap_or(query_embedding);
            match project_store.search_filtered_with_index(
                search_embedding,
                &bridge_filter,
                bridge_limit,
                opts.seed_threshold,
                project_index,
            ) {
                Ok(r) if !r.is_empty() => Some((seed.score, r)),
                Ok(_) => None,
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        ref_seed = %seed.chunk.name,
                        "Bridge search failed for ref seed"
                    );
                    bridge_error_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    None
                }
            }
        })
        .collect();

    // Surface total bridge search failures that silently degrade quality
    let total_bridge_errors = bridge_error_count.load(std::sync::atomic::Ordering::Relaxed);
    if total_bridge_errors > 0 {
        tracing::warn!(
            failed = total_bridge_errors,
            total = ref_seeds.len(),
            "Bridge searches failed — cross-index gather quality may be reduced"
        );
    }

    drop(_bridge_span);

    // Merge into bridge_scores sequentially (HashMap not Sync).
    // Consume `bridge_results` so the per-result strings (`pr.chunk.name`,
    // `pr.chunk.id`) move into the entry instead of being cloned and
    // immediately dropped along with the source vec.
    let mut bridge_scores: HashMap<String, (f32, String)> = HashMap::new(); // name -> (score, chunk_id)
    for (seed_score, results) in bridge_results {
        for pr in results {
            let bridge_score = pr.score * seed_score;
            match bridge_scores.entry(pr.chunk.name) {
                std::collections::hash_map::Entry::Vacant(e) => {
                    e.insert((bridge_score, pr.chunk.id));
                }
                std::collections::hash_map::Entry::Occupied(mut e) => {
                    if bridge_score > e.get().0 {
                        e.insert((bridge_score, pr.chunk.id));
                    }
                }
            }
        }
    }

    tracing::debug!(bridge_count = bridge_scores.len(), "Bridge search complete");

    if bridge_scores.is_empty() {
        // No project code found — return ref seeds only
        let mut result_chunks = ref_chunks;
        result_chunks.truncate(opts.limit);
        return Ok(GatherResult {
            chunks: result_chunks,
            expansion_capped: false,
            search_degraded: false,
        });
    }

    // 4. BFS expand project-side bridges via project call graph
    let graph = project_store.get_call_graph()?;

    let mut name_scores: HashMap<String, (f32, usize)> = HashMap::new();
    for (name, (score, _)) in &bridge_scores {
        name_scores.insert(name.clone(), (*score, 0));
    }

    let expansion_capped = bfs_expand(&mut name_scores, &graph, opts);
    tracing::debug!(
        expanded_nodes = name_scores.len(),
        expansion_capped,
        "Project BFS expansion complete"
    );

    // 5. Batch-fetch project chunks
    // Cross-index gather seeds from the reference, not the project store —
    // the worktree overlay does not apply (Part A), so no overlay seeds.
    let (project_chunks, search_degraded) =
        fetch_and_assemble(project_store, &name_scores, root, None);

    // 6. Combine ref seeds + project chunks, sort by score, truncate, re-sort to reading order
    let mut all_chunks = ref_chunks;
    all_chunks.extend(project_chunks);

    // Sort by score desc with a full tiebreak cascade (name → file → line_start).
    // Name alone isn't unique — two chunks named `new` from different files
    // collide — so we also compare file and line_start to keep the truncate()
    // boundary deterministic across process invocations.
    all_chunks.sort_by(|a, b| {
        b.score
            .total_cmp(&a.score)
            .then(a.name.cmp(&b.name))
            .then(a.file.cmp(&b.file))
            .then(a.line_start.cmp(&b.line_start))
    });
    all_chunks.truncate(opts.limit);
    // Sort: ref chunks first (by source name), then project chunks, each group in file/line order
    all_chunks.sort_by(|a, b| {
        // Reference chunks come first, project chunks second
        let source_ord = match (&a.source, &b.source) {
            (Some(_), None) => std::cmp::Ordering::Less,
            (None, Some(_)) => std::cmp::Ordering::Greater,
            _ => std::cmp::Ordering::Equal,
        };
        source_ord
            .then(a.file.cmp(&b.file))
            .then(a.line_start.cmp(&b.line_start))
            .then(a.name.cmp(&b.name))
    });

    Ok(GatherResult {
        chunks: all_chunks,
        expansion_capped,
        search_degraded,
    })
}

/// Get neighbors in the specified direction.
/// Returns `Arc<str>` cloned from the CallGraph maps to avoid per-node String allocations.
fn get_neighbors(graph: &CallGraph, name: &str, direction: GatherDirection) -> Vec<Arc<str>> {
    let mut neighbors = Vec::new();
    match direction {
        GatherDirection::Callees | GatherDirection::Both => {
            if let Some(callees) = graph.forward.get(name) {
                neighbors.extend(callees.iter().map(Arc::clone));
            }
        }
        _ => {}
    }
    match direction {
        GatherDirection::Callers | GatherDirection::Both => {
            if let Some(callers) = graph.reverse.get(name) {
                neighbors.extend(callers.iter().map(Arc::clone));
            }
        }
        _ => {}
    }
    neighbors
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::assert_matches;

    /// Constructs a sample call graph representing function call relationships.
    ///
    /// Creates a CallGraph with both forward and reverse call mappings. The forward map tracks which functions each function calls, while the reverse map tracks which functions call each function.
    ///
    /// The example graph has the following structure:
    /// - Function A calls B and C
    /// - Function B calls D
    /// - Function C and D are called by A and B respectively
    ///
    /// # Returns
    ///
    /// A CallGraph containing the forward and reverse call mappings for the sample graph.
    fn make_graph() -> CallGraph {
        let mut forward = HashMap::new();
        let mut reverse = HashMap::new();

        // A calls B and C
        forward.insert("A".to_string(), vec!["B".to_string(), "C".to_string()]);
        // B calls D
        forward.insert("B".to_string(), vec!["D".to_string()]);

        // B and C are called by A
        reverse.insert("B".to_string(), vec!["A".to_string()]);
        reverse.insert("C".to_string(), vec!["A".to_string()]);
        // D is called by B
        reverse.insert("D".to_string(), vec!["B".to_string()]);

        CallGraph::from_string_maps(forward, reverse)
    }

    #[test]
    fn test_direction_parse() {
        assert_matches!(
            "both".parse::<GatherDirection>().unwrap(),
            GatherDirection::Both
        );
        assert_matches!(
            "callers".parse::<GatherDirection>().unwrap(),
            GatherDirection::Callers
        );
        assert_matches!(
            "callees".parse::<GatherDirection>().unwrap(),
            GatherDirection::Callees
        );
        assert!("invalid".parse::<GatherDirection>().is_err());
    }

    #[test]
    fn test_default_options() {
        let opts = GatherOptions::default();
        assert_eq!(opts.expand_depth, 1);
        assert_eq!(opts.limit, 10);
        assert_matches!(opts.direction, GatherDirection::Both);
    }

    #[test]
    fn test_get_neighbors_callees() {
        let graph = make_graph();
        let neighbors = get_neighbors(&graph, "A", GatherDirection::Callees);
        assert_eq!(neighbors.len(), 2);
        assert!(neighbors.contains(&Arc::from("B")));
        assert!(neighbors.contains(&Arc::from("C")));
    }

    #[test]
    fn test_get_neighbors_callers() {
        let graph = make_graph();
        let neighbors = get_neighbors(&graph, "B", GatherDirection::Callers);
        assert_eq!(neighbors.len(), 1);
        assert_eq!(&*neighbors[0], "A");
    }

    #[test]
    fn test_get_neighbors_both() {
        let graph = make_graph();
        // B has callees [D] and callers [A]
        let neighbors = get_neighbors(&graph, "B", GatherDirection::Both);
        assert_eq!(neighbors.len(), 2);
        assert!(neighbors.contains(&Arc::from("D")));
        assert!(neighbors.contains(&Arc::from("A")));
    }

    #[test]
    fn test_get_neighbors_unknown_node() {
        let graph = make_graph();
        let neighbors = get_neighbors(&graph, "Z", GatherDirection::Both);
        assert!(neighbors.is_empty());
    }

    #[test]
    fn test_get_neighbors_leaf_node() {
        let graph = make_graph();
        // D has no callees, only callers
        let callees = get_neighbors(&graph, "D", GatherDirection::Callees);
        assert!(callees.is_empty());

        let callers = get_neighbors(&graph, "D", GatherDirection::Callers);
        assert_eq!(callers.len(), 1);
        assert_eq!(&*callers[0], "B");
    }

    /// The seed-sort cascade in `bfs_expand` must be deterministic on
    /// `(score desc, name asc)`. Directly testing `bfs_expand` requires a
    /// `Store`, so we pin the seed-ordering logic that sits at the top of
    /// the function.
    #[test]
    fn test_p2_45_bfs_seed_order_is_deterministic() {
        // Build a name_scores-shaped vector with two equally-scored seeds
        // and one above. Sort 50× from different starting orders;
        // assert every run produces the same sequence.
        let inputs: Vec<Vec<(&str, f32)>> = (0..50)
            .map(|seed| {
                let mut v = vec![("foo", 0.5_f32), ("bar", 0.5_f32), ("zed", 0.9_f32)];
                // Deterministic LCG scramble keyed on `seed`.
                let mut state = (seed as u64)
                    .wrapping_mul(2862933555777941757)
                    .wrapping_add(1);
                for i in (1..v.len()).rev() {
                    state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
                    let j = (state as usize) % (i + 1);
                    v.swap(i, j);
                }
                v
            })
            .collect();
        let canonical: Vec<&str> = {
            let mut v = inputs[0].clone();
            v.sort_by(|(a_name, a_score), (b_name, b_score)| {
                b_score
                    .partial_cmp(a_score)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| a_name.cmp(b_name))
            });
            v.into_iter().map(|(n, _)| n).collect()
        };
        for (i, mut perm) in inputs.into_iter().enumerate() {
            perm.sort_by(|(a_name, a_score), (b_name, b_score)| {
                b_score
                    .partial_cmp(a_score)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| a_name.cmp(b_name))
            });
            let names: Vec<&str> = perm.into_iter().map(|(n, _)| n).collect();
            assert_eq!(
                names, canonical,
                "permutation #{i} produced a different seed order"
            );
        }
        // Spot-check: highest score wins, ties broken by name asc.
        assert_eq!(canonical, vec!["zed", "bar", "foo"]);
    }

    #[test]
    fn test_gather_options_builder() {
        let opts = GatherOptions::default()
            .with_expand_depth(3)
            .with_direction(GatherDirection::Callers)
            .with_limit(20)
            .with_seed_limit(10)
            .with_seed_threshold(0.5)
            .with_decay_factor(0.9);
        assert_eq!(opts.expand_depth, 3);
        assert_matches!(opts.direction, GatherDirection::Callers);
        assert_eq!(opts.limit, 20);
        assert_eq!(opts.seed_limit, 10);
        assert!((opts.seed_threshold - 0.5).abs() < f32::EPSILON);
        assert!((opts.decay_factor - 0.9).abs() < f32::EPSILON);
    }

    #[test]
    fn test_bfs_seed_order_deterministic() {
        // BFS seeding must be deterministic regardless of HashMap iteration
        // order. Enqueueing seeds in `name_scores.keys()` order is
        // process-seed-randomized, and when the cap fires mid-expansion that
        // randomness would bleed into the output. Seed with 10 names, cap at
        // 15, run twice, assert output is the same set on both runs.
        let mut forward = HashMap::new();
        let mut reverse = HashMap::new();
        // Build a graph where every seed has multiple callees so the cap
        // has a real decision to make.
        for seed_idx in 0..10 {
            let seed_name = format!("S{seed_idx}");
            let callees: Vec<String> = (0..5).map(|i| format!("callee_{seed_idx}_{i}")).collect();
            forward.insert(seed_name.clone(), callees.clone());
            for c in &callees {
                reverse
                    .entry(c.clone())
                    .or_insert_with(Vec::new)
                    .push(seed_name.clone());
            }
        }
        let graph = CallGraph::from_string_maps(forward, reverse);

        let build_seeds = || {
            let mut m: HashMap<String, (f32, usize)> = HashMap::new();
            for i in 0..10 {
                // Give all seeds the same score so name-ascending is the
                // only tie-break left — the worst case for HashMap-order
                // leakage.
                m.insert(format!("S{i}"), (0.5, 0));
            }
            m
        };

        let opts = GatherOptions::default()
            .with_expand_depth(1)
            .with_direction(GatherDirection::Callees)
            .with_max_expanded_nodes(15); // 10 seeds + 5 callees = first seed's fanout only

        let mut run1 = build_seeds();
        let capped1 = bfs_expand(&mut run1, &graph, &opts);
        let mut run2 = build_seeds();
        let capped2 = bfs_expand(&mut run2, &graph, &opts);

        assert_eq!(capped1, capped2, "cap flag must match across runs");
        assert!(
            capped1,
            "cap should fire at 15 nodes given 10 seeds x 5 callees"
        );

        // Set of included names must match across runs.
        let mut keys1: Vec<&String> = run1.keys().collect();
        let mut keys2: Vec<&String> = run2.keys().collect();
        keys1.sort();
        keys2.sort();
        assert_eq!(
            keys1, keys2,
            "BFS expanded the same nodes across two runs with identical input"
        );

        // And the first-processed seed (which gets full fanout before the
        // cap fires) should be S0 under (score desc, name asc) sort —
        // not a random HashMap-seed-dependent choice.
        assert!(
            run1.contains_key("callee_0_0"),
            "seed S0 should be processed first under name-ascending sort"
        );
    }

    /// When the cap fires partway through a node's neighbour loop,
    /// already-visited neighbours that *would* have had their score bumped to
    /// a higher value still get the bump: the cap check lives inside the
    /// `!visited` branch (where new insertions actually grow `name_scores`),
    /// not above the visited / bump branches.
    #[test]
    fn bfs_expand_score_bump_runs_when_cap_reached() {
        // Construct a minimal graph that exercises the bump-after-cap
        // branch deterministically:
        //
        //   Seed -> [Target, new_a, new_b]
        //
        // Pre-populate `name_scores` with both Seed (score 1.0, depth 0)
        // AND Target at a deliberately-low score (0.10 at depth 99) so
        // Target is *already-visited* before BFS runs. With cap = 2:
        //   - name_scores starts with {Seed, Target} == 2 entries.
        //   - Seed expands. First neighbour is Target → already-visited
        //     branch. The bump branch runs (no cap check above it), Target
        //     gets promoted from 0.10 to 1.0*0.5=0.5.
        //   - Next neighbour `new_a` → !visited branch → cap check fires
        //     (len == 2 == cap), sets `expansion_capped=true`, breaks.
        let mut forward = HashMap::new();
        let mut reverse = HashMap::new();
        forward.insert(
            "Seed".to_string(),
            vec![
                "Target".to_string(),
                "new_a".to_string(),
                "new_b".to_string(),
            ],
        );
        for callee in ["Target", "new_a", "new_b"] {
            reverse.insert(callee.to_string(), vec!["Seed".to_string()]);
        }
        let graph = CallGraph::from_string_maps(forward, reverse);

        let mut ns: HashMap<String, (f32, usize)> = HashMap::new();
        ns.insert("Seed".to_string(), (1.0, 0));
        ns.insert("Target".to_string(), (0.10, 99));

        let opts = GatherOptions::default()
            .with_expand_depth(1)
            .with_direction(GatherDirection::Callees)
            .with_decay_factor(0.5)
            // Cap == 2: Seed + Target == 2. Any new insertion trips it.
            .with_max_expanded_nodes(2);

        let capped = bfs_expand(&mut ns, &graph, &opts);
        assert!(capped, "cap MUST fire under this configuration");

        // Target's score MUST have been bumped from 0.10 to Seed's decayed
        // score (1.0 * 0.5 = 0.5) — the bump runs even after the cap fires.
        let (target_score, target_depth) = ns["Target"];
        assert!(
            (target_score - 0.5).abs() < f32::EPSILON,
            "Target score should have been bumped to 0.5 (Seed*decay), got {target_score}",
        );
        // Depth bump: existing 99 → min(99, 0+1) = 1.
        assert_eq!(target_depth, 1, "Target depth should be bumped to 1");
    }

    #[test]
    fn test_bfs_depth_preserves_minimum() {
        // Graph: A -> B -> D, A -> C -> D (two paths to D)
        // If A is the seed at depth 0, B and C are discovered at depth 1.
        // D is first reached via B at depth 2, then via C also at depth 2.
        // But if B has a higher score, it should update D's score without
        // overwriting D's depth to a deeper value.
        let mut forward = HashMap::new();
        let mut reverse = HashMap::new();

        forward.insert("A".to_string(), vec!["B".to_string(), "C".to_string()]);
        forward.insert("B".to_string(), vec!["D".to_string()]);
        forward.insert("C".to_string(), vec!["D".to_string()]);

        reverse.insert("B".to_string(), vec!["A".to_string()]);
        reverse.insert("C".to_string(), vec!["A".to_string()]);
        reverse.insert("D".to_string(), vec!["B".to_string(), "C".to_string()]);

        let graph = CallGraph::from_string_maps(forward, reverse);

        let mut name_scores = HashMap::new();
        name_scores.insert("A".to_string(), (1.0, 0));

        let opts = GatherOptions::default()
            .with_expand_depth(3)
            .with_direction(GatherDirection::Callees)
            .with_decay_factor(0.8);

        bfs_expand(&mut name_scores, &graph, &opts);

        // D should be discovered at depth 2 (A->B->D or A->C->D)
        // and should keep the minimum depth even if score is updated
        let (_, depth) = name_scores["D"];
        assert_eq!(depth, 2, "D should preserve minimum depth of 2");

        // B and C should be at depth 1
        assert_eq!(name_scores["B"].1, 1);
        assert_eq!(name_scores["C"].1, 1);
    }
}
