//! Partition-aware ID allocation system.
//!
//! This module implements a composite ID encoding scheme that embeds partition
//! information directly into u64 IDs while maintaining backward compatibility
//! with existing sequential ID allocation.
//!
//! # ID Structure
//!
//! ```text
//! Bits [63-48]: Partition ID (16 bits = 65,536 max partitions)
//! Bits [47-0]:  Local ID (48 bits = 281T IDs per partition)
//!
//! Examples:
//! Partition 0, ID 1:  0x0000_0000_0000_0001 (backward compatible)
//! Partition 5, ID 42: 0x0005_0000_0000_002A
//! ```

use std::fmt;

/// Maximum number of partitions (2^16)
pub const MAX_PARTITIONS: u32 = 65536;

/// Maximum local ID per partition (2^48)
pub const MAX_LOCAL_ID: u64 = 281_474_976_710_656;

/// A partitioned ID that encodes both partition and local ID.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PartitionedId(u64);

impl PartitionedId {
    /// Create a new partitioned ID from partition and local ID components.
    ///
    /// # Panics
    ///
    /// Panics if partition_id >= MAX_PARTITIONS or local_id >= MAX_LOCAL_ID.
    pub fn new(partition_id: u32, local_id: u64) -> Self {
        assert!(
            partition_id < MAX_PARTITIONS,
            "Partition ID {} exceeds maximum {}",
            partition_id,
            MAX_PARTITIONS
        );
        assert!(
            local_id < MAX_LOCAL_ID,
            "Local ID {} exceeds maximum {}",
            local_id,
            MAX_LOCAL_ID
        );

        // Encode: [partition_id (16 bits)] [local_id (48 bits)]
        let composite = ((partition_id as u64) << 48) | (local_id & 0xFFFF_FFFF_FFFF);
        PartitionedId(composite)
    }

    /// Create a partitioned ID from a raw u64 value.
    ///
    /// This is useful for deserialization and backward compatibility.
    pub fn from_raw(raw: u64) -> Self {
        PartitionedId(raw)
    }

    /// Get the raw u64 representation.
    pub fn as_raw(&self) -> u64 {
        self.0
    }

    /// Extract the partition ID (upper 16 bits).
    pub fn partition_id(&self) -> u32 {
        (self.0 >> 48) as u32
    }

    /// Extract the local ID (lower 48 bits).
    pub fn local_id(&self) -> u64 {
        self.0 & 0xFFFF_FFFF_FFFF
    }

    /// Check if this ID belongs to the given partition.
    pub fn is_in_partition(&self, partition_id: u32) -> bool {
        self.partition_id() == partition_id
    }
}

impl fmt::Display for PartitionedId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "PartitionedId(partition={}, local={})",
            self.partition_id(),
            self.local_id()
        )
    }
}

impl From<u64> for PartitionedId {
    fn from(raw: u64) -> Self {
        PartitionedId::from_raw(raw)
    }
}

impl From<PartitionedId> for u64 {
    fn from(id: PartitionedId) -> Self {
        id.as_raw()
    }
}

/// Allocator for partition-local IDs.
///
/// Each partition maintains its own sequential counter for local IDs.
/// This allocator is thread-safe via interior mutability (use Arc<Mutex<...>>
/// for multi-threaded access).
#[derive(Debug, Clone)]
pub struct PartitionIdAllocator {
    partition_id: u32,
    next_local_id: u64,
}

impl PartitionIdAllocator {
    /// Create a new ID allocator for the given partition.
    ///
    /// # Arguments
    ///
    /// * `partition_id` - The partition this allocator belongs to
    /// * `start_id` - The starting local ID (default: 1 for compatibility)
    pub fn new(partition_id: u32, start_id: u64) -> Self {
        assert!(
            partition_id < MAX_PARTITIONS,
            "Partition ID {} exceeds maximum {}",
            partition_id,
            MAX_PARTITIONS
        );
        assert!(
            start_id < MAX_LOCAL_ID,
            "Start ID {} exceeds maximum {}",
            start_id,
            MAX_LOCAL_ID
        );

        Self {
            partition_id,
            next_local_id: start_id,
        }
    }

    /// Allocate the next node ID for this partition.
    pub fn allocate_node_id(&mut self) -> PartitionedId {
        let local_id = self.next_local_id;
        self.next_local_id += 1;

        if self.next_local_id >= MAX_LOCAL_ID {
            panic!(
                "Partition {} has exhausted local ID space ({} IDs allocated)",
                self.partition_id, MAX_LOCAL_ID
            );
        }

        PartitionedId::new(self.partition_id, local_id)
    }

    /// Allocate the next edge ID for this partition.
    ///
    /// Note: This uses the same local ID space as nodes. In the current
    /// implementation, nodes and edges share the same counter. For
    /// separate counters, use a separate allocator instance.
    pub fn allocate_edge_id(&mut self) -> PartitionedId {
        self.allocate_node_id()
    }

    /// Get the partition ID this allocator belongs to.
    pub fn partition_id(&self) -> u32 {
        self.partition_id
    }

    /// Get the next local ID that will be allocated.
    pub fn peek_next_local_id(&self) -> u64 {
        self.next_local_id
    }

    /// Reset the allocator to a specific local ID.
    ///
    /// This is useful for loading partitions from storage where IDs
    /// have already been assigned.
    pub fn reset_to(&mut self, next_local_id: u64) {
        assert!(
            next_local_id < MAX_LOCAL_ID,
            "Next local ID {} exceeds maximum {}",
            next_local_id,
            MAX_LOCAL_ID
        );
        self.next_local_id = next_local_id;
    }
}

impl Default for PartitionIdAllocator {
    /// Create a default allocator for partition 0, starting at ID 1.
    ///
    /// This maintains backward compatibility with single-instance mode.
    fn default() -> Self {
        Self::new(0, 1)
    }
}

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

    #[test]
    fn test_partitioned_id_encoding() {
        // Partition 0, ID 1 (backward compatible)
        let id = PartitionedId::new(0, 1);
        assert_eq!(id.as_raw(), 1);
        assert_eq!(id.partition_id(), 0);
        assert_eq!(id.local_id(), 1);

        // Partition 5, ID 42
        let id = PartitionedId::new(5, 42);
        assert_eq!(id.as_raw(), 0x0005_0000_0000_002A);
        assert_eq!(id.partition_id(), 5);
        assert_eq!(id.local_id(), 42);

        // Maximum partition ID
        let id = PartitionedId::new(65535, 1);
        assert_eq!(id.partition_id(), 65535);
        assert_eq!(id.local_id(), 1);

        // Large local ID
        let id = PartitionedId::new(1, 1_000_000_000_000);
        assert_eq!(id.partition_id(), 1);
        assert_eq!(id.local_id(), 1_000_000_000_000);
    }

    #[test]
    fn test_partitioned_id_is_in_partition() {
        let id = PartitionedId::new(5, 42);
        assert!(id.is_in_partition(5));
        assert!(!id.is_in_partition(0));
        assert!(!id.is_in_partition(6));
    }

    #[test]
    fn test_partitioned_id_from_raw() {
        let raw = 0x0005_0000_0000_002A_u64;
        let id = PartitionedId::from_raw(raw);
        assert_eq!(id.partition_id(), 5);
        assert_eq!(id.local_id(), 42);
    }

    #[test]
    fn test_partitioned_id_conversions() {
        let id = PartitionedId::new(3, 100);
        let raw: u64 = id.into();
        let id2: PartitionedId = raw.into();
        assert_eq!(id, id2);
    }

    #[test]
    fn test_allocator_sequential() {
        let mut allocator = PartitionIdAllocator::new(2, 1);

        let id1 = allocator.allocate_node_id();
        assert_eq!(id1.partition_id(), 2);
        assert_eq!(id1.local_id(), 1);

        let id2 = allocator.allocate_node_id();
        assert_eq!(id2.partition_id(), 2);
        assert_eq!(id2.local_id(), 2);

        let id3 = allocator.allocate_edge_id();
        assert_eq!(id3.partition_id(), 2);
        assert_eq!(id3.local_id(), 3);
    }

    #[test]
    fn test_allocator_default() {
        let mut allocator = PartitionIdAllocator::default();
        let id = allocator.allocate_node_id();
        assert_eq!(id.as_raw(), 1); // Backward compatible
        assert_eq!(id.partition_id(), 0);
        assert_eq!(id.local_id(), 1);
    }

    #[test]
    fn test_allocator_reset() {
        let mut allocator = PartitionIdAllocator::new(1, 1);
        allocator.allocate_node_id(); // ID 1
        allocator.allocate_node_id(); // ID 2

        allocator.reset_to(10);
        let id = allocator.allocate_node_id();
        assert_eq!(id.local_id(), 10);
    }

    #[test]
    #[should_panic(expected = "Partition ID 65536 exceeds maximum")]
    fn test_partition_id_overflow() {
        PartitionedId::new(65536, 1);
    }

    #[test]
    #[should_panic(expected = "Local ID")]
    fn test_local_id_overflow() {
        PartitionedId::new(0, MAX_LOCAL_ID);
    }
}
