//! Graph partitioning by module.
//!
//! Splits large knowledge graphs into per-module partitions based on
//! top-level directory structure. Provides a `PartitionRouter` that
//! routes queries to relevant partitions for parallel execution.

use std::collections::HashMap;
use std::path::Path;
use std::sync::RwLock;

use tracing::debug;

/// A named partition containing file paths.
#[derive(Debug, Clone)]
pub struct Partition {
    /// Partition name (typically the top-level directory).
    pub name: String,
    /// All file paths belonging to this partition.
    pub files: Vec<String>,
}

/// Routes queries to relevant graph partitions.
///
/// Partitions are built by top-level directory: files in `src/auth/` belong
/// to the "src/auth" partition, `tests/` files to "tests", root files to
/// the "(root)" partition.
pub struct PartitionRouter {
    /// Map from partition name → file paths.
    partitions: RwLock<HashMap<String, Vec<String>>>,
    /// Reverse map from file path → partition name.
    file_to_partition: RwLock<HashMap<String, String>>,
    /// Maximum directory depth to partition on.
    max_depth: usize,
}

impl PartitionRouter {
    /// Create a new router that partitions by directories up to `max_depth`.
    pub fn new(max_depth: usize) -> Self {
        Self {
            partitions: RwLock::new(HashMap::new()),
            file_to_partition: RwLock::new(HashMap::new()),
            max_depth: max_depth.max(1),
        }
    }

    /// Build partitions from a list of file paths.
    pub fn build_from_files(&self, file_paths: &[String]) {
        let mut partitions: HashMap<String, Vec<String>> = HashMap::new();
        let mut file_map: HashMap<String, String> = HashMap::new();

        for path in file_paths {
            let partition = self.partition_for_path(path);
            file_map.insert(path.clone(), partition.clone());
            partitions.entry(partition).or_default().push(path.clone());
        }

        let count = partitions.len();
        if let Ok(mut p) = self.partitions.write() {
            *p = partitions;
        }
        if let Ok(mut f) = self.file_to_partition.write() {
            *f = file_map;
        }

        debug!(
            partitions = count,
            files = file_paths.len(),
            "Built graph partitions"
        );
    }

    /// Determine which partition a file path belongs to.
    fn partition_for_path(&self, file_path: &str) -> String {
        let path = Path::new(file_path);
        let components: Vec<&str> = path
            .components()
            .filter_map(|c| {
                if let std::path::Component::Normal(name) = c {
                    name.to_str()
                } else {
                    None
                }
            })
            .collect();

        if components.len() <= 1 {
            return "(root)".to_string();
        }

        // Take up to max_depth directory components (excluding the filename).
        let dir_components = &components[..components.len() - 1];
        let depth = dir_components.len().min(self.max_depth);
        dir_components[..depth].join("/")
    }

    /// Find which partitions are relevant for a query.
    ///
    /// Returns partition names that should be searched. If no partition
    /// hints are available, returns all partitions.
    pub fn route_query(&self, query_hints: &[String]) -> Vec<String> {
        let partitions = match self.partitions.read() {
            Ok(p) => p,
            Err(_) => return vec![],
        };

        if query_hints.is_empty() {
            // No hints — search all partitions.
            return partitions.keys().cloned().collect();
        }

        let file_to_partition = match self.file_to_partition.read() {
            Ok(f) => f,
            Err(_) => return partitions.keys().cloned().collect(),
        };

        let mut relevant: Vec<String> = Vec::new();

        for hint in query_hints {
            // Check if the hint is a file path we know about.
            if let Some(partition) = file_to_partition.get(hint) {
                if !relevant.contains(partition) {
                    relevant.push(partition.clone());
                }
                continue;
            }

            // Check if the hint matches a partition name prefix.
            for partition_name in partitions.keys() {
                if partition_name.contains(hint) && !relevant.contains(partition_name) {
                    relevant.push(partition_name.clone());
                }
            }
        }

        // If no relevant partitions found, return all.
        if relevant.is_empty() {
            partitions.keys().cloned().collect()
        } else {
            relevant
        }
    }

    /// Get all files in a specific partition.
    pub fn files_in_partition(&self, partition_name: &str) -> Vec<String> {
        self.partitions
            .read()
            .ok()
            .and_then(|p| p.get(partition_name).cloned())
            .unwrap_or_default()
    }

    /// Get the partition name for a given file.
    pub fn partition_of(&self, file_path: &str) -> Option<String> {
        self.file_to_partition
            .read()
            .ok()
            .and_then(|f| f.get(file_path).cloned())
    }

    /// Number of partitions.
    pub fn partition_count(&self) -> usize {
        self.partitions.read().map(|p| p.len()).unwrap_or(0)
    }

    /// Get all partition names with their file counts.
    pub fn partition_summary(&self) -> Vec<(String, usize)> {
        self.partitions
            .read()
            .map(|p| {
                let mut summary: Vec<_> = p.iter().map(|(k, v)| (k.clone(), v.len())).collect();
                summary.sort_by(|a, b| b.1.cmp(&a.1)); // Sort by file count descending
                summary
            })
            .unwrap_or_default()
    }
}

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

    fn test_files() -> Vec<String> {
        vec![
            "src/main.rs".into(),
            "src/lib.rs".into(),
            "src/agents/gandalf.rs".into(),
            "src/agents/gimli.rs".into(),
            "src/db/sqlite.rs".into(),
            "src/db/kuzu.rs".into(),
            "src/search/orchestrator.rs".into(),
            "tests/unit_test.rs".into(),
            "tests/integration_test.rs".into(),
            "README.md".into(),
        ]
    }

    #[test]
    fn test_partition_by_top_level() {
        let router = PartitionRouter::new(1);
        router.build_from_files(&test_files());

        assert!(router.partition_count() >= 3); // src, tests, (root)
        assert!(router.files_in_partition("src").len() >= 2);
        assert_eq!(router.files_in_partition("tests").len(), 2);
        assert_eq!(router.files_in_partition("(root)").len(), 1);
    }

    #[test]
    fn test_partition_depth_2() {
        let router = PartitionRouter::new(2);
        router.build_from_files(&test_files());

        // With depth 2, "src/agents" and "src/db" are separate partitions.
        assert!(router.files_in_partition("src/agents").len() >= 2);
        assert!(router.files_in_partition("src/db").len() >= 2);
    }

    #[test]
    fn test_partition_of_file() {
        let router = PartitionRouter::new(2);
        router.build_from_files(&test_files());

        assert_eq!(
            router.partition_of("src/agents/gandalf.rs"),
            Some("src/agents".into())
        );
        assert_eq!(
            router.partition_of("README.md"),
            Some("(root)".into())
        );
    }

    #[test]
    fn test_route_query_with_file_hint() {
        let router = PartitionRouter::new(2);
        router.build_from_files(&test_files());

        let relevant = router.route_query(&["src/agents/gandalf.rs".into()]);
        assert!(relevant.contains(&"src/agents".into()));
    }

    #[test]
    fn test_route_query_with_dir_hint() {
        let router = PartitionRouter::new(2);
        router.build_from_files(&test_files());

        let relevant = router.route_query(&["agents".into()]);
        assert!(relevant.contains(&"src/agents".into()));
    }

    #[test]
    fn test_route_query_no_hints_returns_all() {
        let router = PartitionRouter::new(1);
        router.build_from_files(&test_files());

        let relevant = router.route_query(&[]);
        assert!(relevant.len() >= 3);
    }

    #[test]
    fn test_partition_summary_sorted() {
        let router = PartitionRouter::new(2);
        router.build_from_files(&test_files());

        let summary = router.partition_summary();
        assert!(!summary.is_empty());
        // Should be sorted by count descending.
        for window in summary.windows(2) {
            assert!(window[0].1 >= window[1].1);
        }
    }

    #[test]
    fn test_root_files() {
        let router = PartitionRouter::new(1);
        router.build_from_files(&["Cargo.toml".into(), "README.md".into()]);
        assert_eq!(router.files_in_partition("(root)").len(), 2);
    }

    #[test]
    fn test_empty_files() {
        let router = PartitionRouter::new(1);
        router.build_from_files(&[]);
        assert_eq!(router.partition_count(), 0);
    }
}
