//! 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 anyhow::Result;

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

/// Options for gather operation
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,
}

impl GatherOptions {
    pub fn with_expand_depth(mut self, depth: usize) -> Self {
        self.expand_depth = depth;
        self
    }
    pub fn with_direction(mut self, direction: GatherDirection) -> Self {
        self.direction = direction;
        self
    }
    pub fn with_limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }
    pub fn with_seed_limit(mut self, limit: usize) -> Self {
        self.seed_limit = limit;
        self
    }
    pub fn with_seed_threshold(mut self, threshold: f32) -> Self {
        self.seed_threshold = threshold;
        self
    }
    pub fn with_decay_factor(mut self, factor: f32) -> Self {
        self.decay_factor = factor;
        self
    }
}

impl Default for GatherOptions {
    fn default() -> Self {
        Self {
            expand_depth: 1,
            direction: GatherDirection::Both,
            limit: 10,
            seed_limit: 5,
            seed_threshold: 0.3,
            decay_factor: 0.8,
        }
    }
}

/// Direction of call graph expansion
#[derive(Debug, Clone, Copy)]
pub enum GatherDirection {
    Both,
    Callers,
    Callees,
}

impl std::str::FromStr for GatherDirection {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "both" => Ok(Self::Both),
            "callers" => Ok(Self::Callers),
            "callees" => Ok(Self::Callees),
            _ => anyhow::bail!("Invalid direction '{}'. Valid: both, callers, callees", s),
        }
    }
}

/// A gathered code chunk with context
#[derive(Debug)]
pub struct GatheredChunk {
    pub name: String,
    pub file: PathBuf,
    pub line_start: u32,
    pub line_end: u32,
    pub signature: String,
    pub content: String,
    pub score: f32,
    pub depth: usize,
}

/// Result of a gather operation
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,
}

/// Maximum nodes in BFS expansion to prevent blowup on hub functions
const MAX_EXPANDED_NODES: usize = 200;

/// Gather relevant code chunks for a query
pub fn gather(
    store: &Store,
    query_embedding: &crate::Embedding,
    query_text: &str,
    opts: &GatherOptions,
    project_root: &Path,
) -> Result<GatherResult> {
    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: true,
        ..SearchFilter::default()
    };
    let seed_results = store.search_filtered(
        query_embedding,
        &filter,
        opts.seed_limit,
        opts.seed_threshold,
    )?;
    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,
        });
    }

    // 2. Load call graph for expansion
    // NOTE: This loads the entire function_calls table into memory each call.
    // Current callers invoke gather() once per request so this is fine.
    // If gather() is ever called in a loop, accept a pre-loaded &CallGraph parameter
    // to avoid redundant loads.
    let graph = store.get_call_graph()?;

    // Seed names with their scores
    let mut name_scores: HashMap<String, (f32, usize)> = HashMap::new(); // name -> (score, depth)
    for r in &seed_results {
        name_scores.insert(r.chunk.name.clone(), (r.score, 0));
    }

    // 3. BFS expand
    let mut expansion_capped = false;
    if opts.expand_depth > 0 {
        let mut queue: VecDeque<(String, usize)> = VecDeque::new();
        for r in &seed_results {
            queue.push_back((r.chunk.name.clone(), 0));
        }

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

            let neighbors = get_neighbors(&graph, &name, opts.direction);
            let base_score = name_scores.get(&name).map(|(s, _)| *s).unwrap_or(0.5);
            let decay = opts.decay_factor.powi((depth + 1) as i32);
            let new_score = base_score * decay;
            for neighbor in neighbors {
                match name_scores.entry(neighbor.clone()) {
                    std::collections::hash_map::Entry::Vacant(e) => {
                        e.insert((new_score, depth + 1));
                        queue.push_back((neighbor, depth + 1));
                    }
                    std::collections::hash_map::Entry::Occupied(mut e) => {
                        if new_score > e.get().0 {
                            e.insert((new_score, depth + 1));
                            // Don't re-add to queue — already explored or queued
                        }
                    }
                }
            }
        }
    }
    tracing::debug!(
        expanded_nodes = name_scores.len(),
        expansion_capped,
        "BFS expansion complete"
    );

    // 4. Batch-fetch chunks for all expanded names, deduplicate by id
    let all_names: Vec<&str> = name_scores.keys().map(|s| s.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 {
        if let Some(results) = batch_results.get(name) {
            if let Some(r) = results.first() {
                // Dedup by chunk id
                if seen_ids.contains(&r.chunk.id) {
                    continue;
                }
                seen_ids.insert(r.chunk.id.clone());

                chunks.push(GatheredChunk {
                    name: r.chunk.name.clone(),
                    file: r
                        .chunk
                        .file
                        .strip_prefix(project_root)
                        .unwrap_or(&r.chunk.file)
                        .to_path_buf(),
                    line_start: r.chunk.line_start,
                    line_end: r.chunk.line_end,
                    signature: r.chunk.signature.clone(),
                    content: r.chunk.content.clone(),
                    score: *score,
                    depth: *depth,
                });
            }
        }
    }

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

    // 5. Sort by score desc (with name tiebreak for determinism), truncate to limit,
    //    then re-sort by file → line_start → name for stable reading order.
    chunks.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(a.name.cmp(&b.name))
    });
    chunks.truncate(opts.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))
    });

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

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

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

    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 { 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(&"B".to_string()));
        assert!(neighbors.contains(&"C".to_string()));
    }

    #[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(&"D".to_string()));
        assert!(neighbors.contains(&"A".to_string()));
    }

    #[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");
    }

    #[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);
    }
}
