//! Partition metadata and mapping structures.
//!
//! This module defines data structures for tracking partition assignments
//! and quality metrics.

use std::collections::{HashMap, HashSet};
use serde::{Deserialize, Serialize};

/// Metadata about a single partition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMetadata {
    /// Partition ID
    pub partition_id: u32,

    /// Number of nodes in this partition
    pub node_count: usize,

    /// Number of edges in this partition
    pub edge_count: usize,

    /// Number of edges with endpoints in other partitions (cross-partition edges)
    pub cross_partition_edge_count: usize,

    /// Set of node labels present in this partition
    pub labels: HashSet<String>,

    /// Set of relationship types present in this partition
    pub relationship_types: HashSet<String>,

    /// Estimated memory usage in bytes
    pub memory_bytes: usize,
}

impl PartitionMetadata {
    /// Create new empty partition metadata.
    pub fn new(partition_id: u32) -> Self {
        Self {
            partition_id,
            node_count: 0,
            edge_count: 0,
            cross_partition_edge_count: 0,
            labels: HashSet::new(),
            relationship_types: HashSet::new(),
            memory_bytes: 0,
        }
    }

    /// Calculate the edge cut ratio (cross-partition edges / total edges).
    pub fn edge_cut_ratio(&self) -> f64 {
        if self.edge_count == 0 {
            0.0
        } else {
            self.cross_partition_edge_count as f64 / self.edge_count as f64
        }
    }
}

/// Maps vertices and edges to their partition assignments.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMap {
    /// Total number of partitions
    pub num_partitions: u32,

    /// Maps vertex ID to partition ID(s). Most vertices belong to exactly one
    /// partition, but boundary vertices may be replicated across partitions
    /// for cross-partition edge handling.
    pub vertex_to_partition: HashMap<u64, u32>,

    /// Maps edge ID to the partition that owns it.
    /// Edges are assigned to the partition of their source vertex.
    pub edge_to_partition: HashMap<u64, u32>,

    /// Metadata for each partition
    pub partition_metadata: Vec<PartitionMetadata>,
}

impl PartitionMap {
    /// Create a new partition map for the specified number of partitions.
    pub fn new(num_partitions: u32) -> Self {
        let mut partition_metadata = Vec::with_capacity(num_partitions as usize);
        for i in 0..num_partitions {
            partition_metadata.push(PartitionMetadata::new(i));
        }

        Self {
            num_partitions,
            vertex_to_partition: HashMap::new(),
            edge_to_partition: HashMap::new(),
            partition_metadata,
        }
    }

    /// Get the partition ID for a vertex.
    pub fn get_vertex_partition(&self, vertex_id: u64) -> Option<u32> {
        self.vertex_to_partition.get(&vertex_id).copied()
    }

    /// Get the partition ID for an edge.
    pub fn get_edge_partition(&self, edge_id: u64) -> Option<u32> {
        self.edge_to_partition.get(&edge_id).copied()
    }

    /// Assign a vertex to a partition.
    pub fn assign_vertex(&mut self, vertex_id: u64, partition_id: u32) {
        assert!(
            partition_id < self.num_partitions,
            "Partition ID {} out of range (max: {})",
            partition_id,
            self.num_partitions - 1
        );

        self.vertex_to_partition.insert(vertex_id, partition_id);
        self.partition_metadata[partition_id as usize].node_count += 1;
    }

    /// Assign an edge to a partition.
    pub fn assign_edge(&mut self, edge_id: u64, partition_id: u32) {
        assert!(
            partition_id < self.num_partitions,
            "Partition ID {} out of range (max: {})",
            partition_id,
            self.num_partitions - 1
        );

        self.edge_to_partition.insert(edge_id, partition_id);
        self.partition_metadata[partition_id as usize].edge_count += 1;
    }

    /// Mark an edge as cross-partition (for quality metrics).
    pub fn mark_cross_partition_edge(&mut self, edge_partition_id: u32) {
        self.partition_metadata[edge_partition_id as usize].cross_partition_edge_count += 1;
    }

    /// Add a label to a partition's metadata.
    pub fn add_label(&mut self, partition_id: u32, label: String) {
        self.partition_metadata[partition_id as usize].labels.insert(label);
    }

    /// Add a relationship type to a partition's metadata.
    pub fn add_relationship_type(&mut self, partition_id: u32, rel_type: String) {
        self.partition_metadata[partition_id as usize].relationship_types.insert(rel_type);
    }

    /// Get metadata for a specific partition.
    pub fn get_partition_metadata(&self, partition_id: u32) -> Option<&PartitionMetadata> {
        self.partition_metadata.get(partition_id as usize)
    }

    /// Calculate overall partition quality metrics.
    pub fn calculate_quality_metrics(&self) -> PartitionQualityMetrics {
        let total_nodes: usize = self.partition_metadata.iter().map(|m| m.node_count).sum();
        let total_edges: usize = self.partition_metadata.iter().map(|m| m.edge_count).sum();
        let total_cross_partition_edges: usize = self.partition_metadata.iter()
            .map(|m| m.cross_partition_edge_count)
            .sum();

        let avg_nodes = if self.num_partitions > 0 {
            total_nodes as f64 / self.num_partitions as f64
        } else {
            0.0
        };

        let max_nodes = self.partition_metadata.iter()
            .map(|m| m.node_count)
            .max()
            .unwrap_or(0);

        let min_nodes = self.partition_metadata.iter()
            .map(|m| m.node_count)
            .min()
            .unwrap_or(0);

        let imbalance = if avg_nodes > 0.0 {
            (max_nodes as f64 - avg_nodes) / avg_nodes
        } else {
            0.0
        };

        let edge_cut_ratio = if total_edges > 0 {
            total_cross_partition_edges as f64 / total_edges as f64
        } else {
            0.0
        };

        PartitionQualityMetrics {
            total_nodes,
            total_edges,
            total_cross_partition_edges,
            avg_nodes_per_partition: avg_nodes,
            max_nodes_per_partition: max_nodes,
            min_nodes_per_partition: min_nodes,
            imbalance,
            edge_cut_ratio,
        }
    }
}

/// Quality metrics for partition assignment.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionQualityMetrics {
    /// Total number of nodes across all partitions
    pub total_nodes: usize,

    /// Total number of edges across all partitions
    pub total_edges: usize,

    /// Total number of edges that cross partition boundaries
    pub total_cross_partition_edges: usize,

    /// Average nodes per partition
    pub avg_nodes_per_partition: f64,

    /// Maximum nodes in any partition
    pub max_nodes_per_partition: usize,

    /// Minimum nodes in any partition
    pub min_nodes_per_partition: usize,

    /// Imbalance factor: (max - avg) / avg
    /// Values closer to 0 are better (more balanced)
    pub imbalance: f64,

    /// Ratio of cross-partition edges to total edges
    /// Values closer to 0 are better (fewer cross-partition queries)
    pub edge_cut_ratio: f64,
}

impl PartitionQualityMetrics {
    /// Check if partitions need rebalancing based on imbalance threshold.
    ///
    /// Returns true if imbalance exceeds 20% (0.2).
    pub fn needs_rebalancing(&self) -> bool {
        self.imbalance > 0.2
    }
}

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

    #[test]
    fn test_partition_map_creation() {
        let map = PartitionMap::new(3);
        assert_eq!(map.num_partitions, 3);
        assert_eq!(map.partition_metadata.len(), 3);
    }

    #[test]
    fn test_vertex_assignment() {
        let mut map = PartitionMap::new(3);

        map.assign_vertex(1, 0);
        map.assign_vertex(2, 1);
        map.assign_vertex(3, 2);

        assert_eq!(map.get_vertex_partition(1), Some(0));
        assert_eq!(map.get_vertex_partition(2), Some(1));
        assert_eq!(map.get_vertex_partition(3), Some(2));

        assert_eq!(map.partition_metadata[0].node_count, 1);
        assert_eq!(map.partition_metadata[1].node_count, 1);
        assert_eq!(map.partition_metadata[2].node_count, 1);
    }

    #[test]
    fn test_edge_assignment() {
        let mut map = PartitionMap::new(2);

        map.assign_edge(100, 0);
        map.assign_edge(101, 1);

        assert_eq!(map.get_edge_partition(100), Some(0));
        assert_eq!(map.get_edge_partition(101), Some(1));

        assert_eq!(map.partition_metadata[0].edge_count, 1);
        assert_eq!(map.partition_metadata[1].edge_count, 1);
    }

    #[test]
    fn test_cross_partition_edges() {
        let mut map = PartitionMap::new(2);

        map.assign_edge(100, 0);
        map.mark_cross_partition_edge(0);

        assert_eq!(map.partition_metadata[0].cross_partition_edge_count, 1);
    }

    #[test]
    fn test_quality_metrics() {
        let mut map = PartitionMap::new(2);

        // Balanced partitions
        map.assign_vertex(1, 0);
        map.assign_vertex(2, 0);
        map.assign_vertex(3, 1);
        map.assign_vertex(4, 1);

        map.assign_edge(100, 0);
        map.assign_edge(101, 1);

        let metrics = map.calculate_quality_metrics();

        assert_eq!(metrics.total_nodes, 4);
        assert_eq!(metrics.total_edges, 2);
        assert_eq!(metrics.avg_nodes_per_partition, 2.0);
        assert_eq!(metrics.max_nodes_per_partition, 2);
        assert_eq!(metrics.min_nodes_per_partition, 2);
        assert_eq!(metrics.imbalance, 0.0); // Perfectly balanced
    }

    #[test]
    fn test_imbalanced_partitions() {
        let mut map = PartitionMap::new(2);

        // Imbalanced: 3 nodes in partition 0, 1 in partition 1
        map.assign_vertex(1, 0);
        map.assign_vertex(2, 0);
        map.assign_vertex(3, 0);
        map.assign_vertex(4, 1);

        let metrics = map.calculate_quality_metrics();

        assert_eq!(metrics.total_nodes, 4);
        assert_eq!(metrics.avg_nodes_per_partition, 2.0);
        assert_eq!(metrics.max_nodes_per_partition, 3);
        assert_eq!(metrics.min_nodes_per_partition, 1);
        assert_eq!(metrics.imbalance, 0.5); // 50% imbalance
        assert!(metrics.needs_rebalancing());
    }

    #[test]
    fn test_edge_cut_ratio() {
        let mut map = PartitionMap::new(2);

        map.assign_edge(100, 0);
        map.assign_edge(101, 0);
        map.assign_edge(102, 1);
        map.assign_edge(103, 1);

        // Mark 2 out of 4 edges as cross-partition
        map.mark_cross_partition_edge(0);
        map.mark_cross_partition_edge(1);

        let metrics = map.calculate_quality_metrics();

        assert_eq!(metrics.total_edges, 4);
        assert_eq!(metrics.total_cross_partition_edges, 2);
        assert_eq!(metrics.edge_cut_ratio, 0.5);
    }
}
