//! Range-based graph partitioning.
//!
//! This module implements a range-based partitioning strategy that assigns
//! contiguous ID ranges to partitions.

use crate::distributed::partition::metadata::PartitionMap;
use crate::distributed::partition::partitioner::{GraphPartitioner, PartitionConfig};
use crate::error::CypherResult;
use crate::graph::PropertyGraph;

/// Range-based partitioner.
///
/// This partitioner assigns vertices to partitions based on contiguous ID ranges:
/// - Partition 0: IDs [min, min + range_size)
/// - Partition 1: IDs [min + range_size, min + 2*range_size)
/// - etc.
///
/// **Pros:**
/// - Simple and fast (O(V + E))
/// - Locality-aware (adjacent IDs likely in same partition)
/// - Good for temporal graphs (newer nodes together)
///
/// **Cons:**
/// - Imbalanced if ID allocation is skewed
/// - Not structure-aware
#[derive(Debug, Default)]
pub struct RangePartitioner;

impl RangePartitioner {
    /// Create a new range partitioner.
    pub fn new() -> Self {
        Self
    }

    /// Determine partition ranges based on vertex IDs.
    fn calculate_ranges(&self, graph: &PropertyGraph, num_partitions: u32) -> Vec<(u64, u64)> {
        let mut vertex_ids: Vec<u64> = graph.all_nodes().map(|n| n.id).collect();

        if vertex_ids.is_empty() {
            return Vec::new();
        }

        vertex_ids.sort_unstable();

        let min_id = vertex_ids[0];
        let max_id = vertex_ids[vertex_ids.len() - 1];
        let range = max_id - min_id + 1;
        let partition_size = (range + num_partitions as u64 - 1) / num_partitions as u64; // Ceiling division

        let mut ranges = Vec::new();
        for i in 0..num_partitions {
            let start = min_id + i as u64 * partition_size;
            let end = if i == num_partitions - 1 {
                max_id + 1 // Last partition gets all remaining IDs
            } else {
                start + partition_size
            };
            ranges.push((start, end));
        }

        ranges
    }

    /// Assign a vertex to a partition based on its ID.
    fn id_to_partition(&self, vertex_id: u64, ranges: &[(u64, u64)]) -> u32 {
        for (i, (start, end)) in ranges.iter().enumerate() {
            if vertex_id >= *start && vertex_id < *end {
                return i as u32;
            }
        }
        // Fallback: assign to last partition
        (ranges.len() - 1) as u32
    }
}

impl GraphPartitioner for RangePartitioner {
    fn partition(&self, graph: &PropertyGraph, config: &PartitionConfig) -> CypherResult<PartitionMap> {
        let mut partition_map = PartitionMap::new(config.num_partitions);

        // Calculate ID ranges for each partition
        let ranges = self.calculate_ranges(graph, config.num_partitions);

        if ranges.is_empty() {
            return Ok(partition_map); // Empty graph
        }

        // Assign vertices to partitions based on ID ranges
        for node in graph.all_nodes() {
            let vertex_id = node.id;
            let partition_id = self.id_to_partition(vertex_id, &ranges);

            partition_map.assign_vertex(vertex_id, partition_id);

            // Track labels
            for label in &node.labels {
                partition_map.add_label(partition_id, label.clone());
            }
        }

        // Assign edges to partitions based on source vertex
        // We need to iterate over all nodes and their edges to get src/dst info
        for node in graph.all_nodes() {
            let src_id = node.id;
            for (from_id, to_id, edge) in graph.all_edges_for_node(src_id) {
                // Only process each edge once (from source perspective)
                if from_id != src_id {
                    continue;
                }

                let src_partition = partition_map.get_vertex_partition(from_id)
                    .expect("Source vertex not assigned to partition");

                partition_map.assign_edge(edge.id, src_partition);

                // Check if edge crosses partition boundaries
                let dst_partition = partition_map.get_vertex_partition(to_id)
                    .expect("Destination vertex not assigned to partition");

                if src_partition != dst_partition {
                    partition_map.mark_cross_partition_edge(src_partition);
                }

                // Track relationship type
                partition_map.add_relationship_type(src_partition, edge.rel_type.clone());
            }
        }

        Ok(partition_map)
    }

    fn strategy_name(&self) -> &str {
        "range"
    }
}

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

    #[test]
    fn test_range_partitioner() {
        let mut graph = PropertyGraph::new();

        // Create 100 nodes with sequential IDs
        execute(&mut graph, "UNWIND range(1, 100) AS i CREATE (n:Node {id: i})").unwrap();

        let partitioner = RangePartitioner::new();
        let config = PartitionConfig::new(4);
        let partition_map = partitioner.partition(&graph, &config).unwrap();

        // Verify all nodes are assigned
        assert_eq!(partition_map.vertex_to_partition.len(), 100);

        // Range partitioning should create reasonably balanced partitions
        let metrics = partition_map.calculate_quality_metrics();
        println!("Range partitioner metrics: {:?}", metrics);
        assert!(metrics.imbalance < 0.3, "Imbalance too high: {}", metrics.imbalance);

        // Verify all partitions have at least some nodes
        for metadata in &partition_map.partition_metadata {
            assert!(metadata.node_count > 0, "Partition {} is empty", metadata.partition_id);
        }
    }

    #[test]
    fn test_range_partitioner_locality() {
        let mut graph = PropertyGraph::new();

        // Create a chain with sequential IDs: 1 -> 2 -> 3 -> 4 -> 5 -> 6
        execute(&mut graph, "
            CREATE (n1:Node {id: 1})-[:NEXT]->(n2:Node {id: 2})
                  -[:NEXT]->(n3:Node {id: 3})-[:NEXT]->(n4:Node {id: 4})
                  -[:NEXT]->(n5:Node {id: 5})-[:NEXT]->(n6:Node {id: 6})
        ").unwrap();

        let partitioner = RangePartitioner::new();
        let config = PartitionConfig::new(2);
        let partition_map = partitioner.partition(&graph, &config).unwrap();

        // Calculate edge cut
        let metrics = partition_map.calculate_quality_metrics();
        println!("Range edge cut ratio: {}", metrics.edge_cut_ratio);

        // Range partitioning should have lower edge cut than hash for sequential graphs
        // (though this depends on the specific graph structure)
    }

    #[test]
    fn test_range_partitioner_empty_graph() {
        let graph = PropertyGraph::new();

        let partitioner = RangePartitioner::new();
        let config = PartitionConfig::new(3);
        let partition_map = partitioner.partition(&graph, &config).unwrap();

        // Empty graph should have empty partition map
        assert_eq!(partition_map.vertex_to_partition.len(), 0);
        assert_eq!(partition_map.edge_to_partition.len(), 0);
    }

    #[test]
    fn test_range_calculate_ranges() {
        let mut graph = PropertyGraph::new();
        // Create 4 nodes - they will get sequential IDs 1, 2, 3, 4
        execute(&mut graph, "
            CREATE (:Node), (:Node), (:Node), (:Node)
        ").unwrap();

        let partitioner = RangePartitioner::new();
        let ranges = partitioner.calculate_ranges(&graph, 2);

        assert_eq!(ranges.len(), 2);
        // Range should split [1, 4] into two partitions
        assert_eq!(ranges[0].0, 1); // First partition starts at 1
        assert!(ranges[1].1 >= 4); // Last partition includes 4
    }
}
