//! Partition analyzer — determines which partitions a query touches.
//!
//! Performs lightweight static analysis of the Cypher query to route it to
//! the minimum set of partitions, reducing cross-partition overhead.
//!
//! ## Analysis rules (in priority order)
//!
//! 1. Query references a node by its full composite ID → single partition.
//! 2. Query references a label that exists on exactly one partition → single partition.
//! 3. Query references labels present on a known subset → multi-partition.
//! 4. Default → all partitions (fan-out).

#[cfg(feature = "distributed")]
use std::collections::{HashMap, HashSet};

#[cfg(feature = "distributed")]
use crate::distributed::partition::PartitionMap;

// ── Distribution type ─────────────────────────────────────────────────────────

/// Which partitions a query needs to touch.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, PartialEq)]
pub enum QueryDistribution {
    /// Only one partition is needed.
    SinglePartition(u32),
    /// A specific subset of partitions.
    MultiPartition(Vec<u32>),
    /// All partitions must be queried (fan-out).
    AllPartitions,
}

// ── Analyzer ──────────────────────────────────────────────────────────────────

/// Stateless query analyzer.
#[cfg(feature = "distributed")]
pub struct PartitionAnalyzer {
    /// label → partition IDs that contain nodes with that label.
    label_index: HashMap<String, HashSet<u32>>,
}

#[cfg(feature = "distributed")]
impl PartitionAnalyzer {
    /// Build the analyzer from the global partition map.
    pub fn from_partition_map(partition_map: &PartitionMap) -> Self {
        let mut label_index: HashMap<String, HashSet<u32>> = HashMap::new();
        for metadata in &partition_map.partition_metadata {
            for label in &metadata.labels {
                label_index
                    .entry(label.clone())
                    .or_default()
                    .insert(metadata.partition_id);
            }
        }
        Self { label_index }
    }

    /// Analyse a Cypher query string and return which partitions are needed.
    ///
    /// This is a best-effort static analysis; it may over-approximate (return
    /// more partitions than strictly needed) but never under-approximate.
    pub fn analyze(&self, query: &str, partition_map: &PartitionMap) -> QueryDistribution {
        // ── Rule 1: ID lookup ─────────────────────────────────────────────
        // Pattern: WHERE id(n) = <number>  or  elementId(n) = <number>
        if let Some(pid) = self.extract_partition_from_id_lookup(query, partition_map) {
            return QueryDistribution::SinglePartition(pid);
        }

        // ── Rule 2 & 3: Label-based routing ───────────────────────────────
        let labels = self.extract_labels_from_query(query);
        if !labels.is_empty() {
            let mut involved: HashSet<u32> = HashSet::new();
            for label in &labels {
                if let Some(partitions) = self.label_index.get(label) {
                    involved.extend(partitions);
                } else {
                    // Unknown label — could be on any partition
                    return QueryDistribution::AllPartitions;
                }
            }
            if !involved.is_empty() {
                let mut ids: Vec<u32> = involved.into_iter().collect();
                ids.sort_unstable();
                return if ids.len() == 1 {
                    QueryDistribution::SinglePartition(ids[0])
                } else {
                    QueryDistribution::MultiPartition(ids)
                };
            }
        }

        // ── Rule 4: Default fan-out ───────────────────────────────────────
        QueryDistribution::AllPartitions
    }

    // ── Private helpers ───────────────────────────────────────────────────

    /// Try to determine the partition from a numeric ID in the query.
    fn extract_partition_from_id_lookup(
        &self,
        query: &str,
        partition_map: &PartitionMap,
    ) -> Option<u32> {
        // Look for patterns like `id(n) = 12345` or `WHERE n.id = 12345`
        let lower = query.to_lowercase();
        let id_val = Self::extract_numeric_after(query, &lower, "id(n) = ")
            .or_else(|| Self::extract_numeric_after(query, &lower, "id(n)="));

        if let Some(node_id) = id_val {
            return partition_map.get_vertex_partition(node_id);
        }
        None
    }

    fn extract_numeric_after(query: &str, lower: &str, pattern: &str) -> Option<u64> {
        let start = lower.find(pattern)?;
        let after = &query[start + pattern.len()..];
        let end = after.find(|c: char| !c.is_ascii_digit()).unwrap_or(after.len());
        after[..end].parse().ok()
    }

    /// Extract node labels from `(n:Label)` patterns, ignoring relationship
    /// types inside `[r:TYPE]` brackets.
    fn extract_labels_from_query(&self, query: &str) -> Vec<String> {
        let mut labels   = Vec::new();
        let mut in_paren: i32 = 0;
        let mut in_bracket: i32 = 0;
        let mut chars = query.chars().peekable();
        while let Some(c) = chars.next() {
            match c {
                '(' => in_paren   += 1,
                ')' => in_paren    = (in_paren - 1).max(0),
                '[' => in_bracket += 1,
                ']' => in_bracket  = (in_bracket - 1).max(0),
                ':' if in_paren > 0 && in_bracket == 0 => {
                    // Only collect labels inside node patterns, not relationships
                    let label: String = chars
                        .by_ref()
                        .take_while(|&ch| ch.is_alphanumeric() || ch == '_')
                        .collect();
                    if !label.is_empty() {
                        labels.push(label);
                    }
                }
                _ => {}
            }
        }
        labels
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;

    fn make_partition_map() -> PartitionMap {
        let mut pm = PartitionMap::new(3);

        // Manually populate partition_metadata labels
        pm.add_label(0, "Person".to_string());
        pm.add_label(0, "Employee".to_string());
        pm.add_relationship_type(0, "KNOWS".to_string());

        pm.add_label(1, "Company".to_string());
        pm.add_relationship_type(1, "WORKS_AT".to_string());

        pm.add_label(2, "Product".to_string());
        pm.add_relationship_type(2, "SELLS".to_string());

        pm
    }

    fn make_analyzer(pm: &PartitionMap) -> PartitionAnalyzer {
        PartitionAnalyzer::from_partition_map(pm)
    }

    #[test]
    fn test_label_routes_to_single_partition() {
        let pm       = make_partition_map();
        let analyzer = make_analyzer(&pm);
        let dist     = analyzer.analyze("MATCH (n:Company) RETURN n", &pm);
        assert_eq!(dist, QueryDistribution::SinglePartition(1));
    }

    #[test]
    fn test_multi_label_routes_to_multiple_partitions() {
        let pm       = make_partition_map();
        let analyzer = make_analyzer(&pm);
        // Person is on partition 0, Company on partition 1
        let dist = analyzer.analyze("MATCH (p:Person)-[:KNOWS]->(c:Company) RETURN p, c", &pm);
        assert_eq!(dist, QueryDistribution::MultiPartition(vec![0, 1]));
    }

    #[test]
    fn test_unknown_label_defaults_to_all() {
        let pm       = make_partition_map();
        let analyzer = make_analyzer(&pm);
        let dist     = analyzer.analyze("MATCH (n:UnknownLabel) RETURN n", &pm);
        assert_eq!(dist, QueryDistribution::AllPartitions);
    }

    #[test]
    fn test_no_label_defaults_to_all() {
        let pm       = make_partition_map();
        let analyzer = make_analyzer(&pm);
        let dist     = analyzer.analyze("MATCH (n) RETURN n", &pm);
        assert_eq!(dist, QueryDistribution::AllPartitions);
    }

    #[test]
    fn test_same_label_multiple_partitions() {
        let mut pm = PartitionMap::new(2);
        pm.add_label(0, "Person".to_string());
        pm.add_label(1, "Person".to_string()); // also has Person

        let analyzer = make_analyzer(&pm);
        let dist     = analyzer.analyze("MATCH (n:Person) RETURN n", &pm);
        // Person is on both partitions → MultiPartition
        assert_eq!(dist, QueryDistribution::MultiPartition(vec![0, 1]));
    }
}
