//! Sharding type definitions for data distribution.
//!
//! This module defines the core types for distributing data across cluster nodes:
//!
//! - [`ShardId`]: Identifier for a data shard
//! - [`RangeDescriptor`]: Describes a shard's key range and replicas
//! - [`MetaRange`]: Collection of all range descriptors (cluster metadata)

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

use crate::membership::NodeId;
use stemedb_core::types::HlcTimestamp;

/// Identifier for a data shard.
///
/// Shards are numbered from 0 to num_shards-1. The mapping from subject
/// to shard is done via consistent hashing (jump hash).
pub type ShardId = u32;

/// Describes a shard's key range, replicas, and metadata.
///
/// Each shard covers a contiguous range of the key space. When shards
/// split or merge, their descriptors are updated atomically in the
/// meta-range.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RangeDescriptor {
    /// Unique identifier for this shard.
    pub shard_id: ShardId,

    /// Start of the key range (inclusive).
    ///
    /// `None` means the range starts at the minimum possible key.
    pub start_key: Option<Vec<u8>>,

    /// End of the key range (exclusive).
    ///
    /// `None` means the range extends to the maximum possible key.
    pub end_key: Option<Vec<u8>>,

    /// Ordered list of replica nodes.
    ///
    /// First node is the leader, subsequent nodes are followers.
    /// Length should equal the replication factor from config.
    pub replicas: Vec<NodeId>,

    /// Current size of data in this shard (bytes).
    ///
    /// Used to trigger split/merge decisions.
    pub size_bytes: u64,

    /// Number of assertions in this shard.
    pub assertion_count: u64,

    /// When this descriptor was last updated (NTP64 time + node_id bytes).
    /// Stored as tuple for serde compatibility.
    #[serde(with = "hlc_serde")]
    pub updated_at: HlcTimestamp,

    /// Generation number for optimistic concurrency.
    ///
    /// Incremented on each update. Used to detect stale reads.
    pub generation: u64,
}

/// Custom serde for HlcTimestamp.
mod hlc_serde {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use stemedb_core::types::HlcTimestamp;

    #[derive(Serialize, Deserialize)]
    struct HlcRepr {
        time_ntp64: u64,
        node_id: [u8; 16],
    }

    pub fn serialize<S>(ts: &HlcTimestamp, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let repr = HlcRepr { time_ntp64: ts.time_ntp64, node_id: ts.node_id };
        repr.serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<HlcTimestamp, D::Error>
    where
        D: Deserializer<'de>,
    {
        let repr = HlcRepr::deserialize(deserializer)?;
        Ok(HlcTimestamp::new(repr.time_ntp64, repr.node_id))
    }
}

impl RangeDescriptor {
    /// Creates a new range descriptor for a full range shard.
    #[must_use]
    pub fn new_full_range(shard_id: ShardId, replicas: Vec<NodeId>) -> Self {
        Self {
            shard_id,
            start_key: None,
            end_key: None,
            replicas,
            size_bytes: 0,
            assertion_count: 0,
            updated_at: HlcTimestamp::default(),
            generation: 1,
        }
    }

    /// Creates a new range descriptor with specific key bounds.
    #[must_use]
    pub fn new(
        shard_id: ShardId,
        start_key: Option<Vec<u8>>,
        end_key: Option<Vec<u8>>,
        replicas: Vec<NodeId>,
    ) -> Self {
        Self {
            shard_id,
            start_key,
            end_key,
            replicas,
            size_bytes: 0,
            assertion_count: 0,
            updated_at: HlcTimestamp::default(),
            generation: 1,
        }
    }

    /// Returns the leader node for this shard.
    #[must_use]
    pub fn leader(&self) -> Option<NodeId> {
        self.replicas.first().copied()
    }

    /// Returns the follower nodes for this shard.
    #[must_use]
    pub fn followers(&self) -> &[NodeId] {
        if self.replicas.len() > 1 {
            &self.replicas[1..]
        } else {
            &[]
        }
    }

    /// Checks if this shard contains the given key.
    #[must_use]
    pub fn contains_key(&self, key: &[u8]) -> bool {
        let after_start =
            self.start_key.as_ref().map(|start| key >= start.as_slice()).unwrap_or(true);

        let before_end = self.end_key.as_ref().map(|end| key < end.as_slice()).unwrap_or(true);

        after_start && before_end
    }

    /// Checks if this shard should be split based on size threshold.
    #[must_use]
    pub fn should_split(&self, threshold_bytes: u64) -> bool {
        self.size_bytes > threshold_bytes
    }

    /// Updates size and assertion count, incrementing generation.
    pub fn update_stats(&mut self, size_bytes: u64, assertion_count: u64, timestamp: HlcTimestamp) {
        self.size_bytes = size_bytes;
        self.assertion_count = assertion_count;
        self.updated_at = timestamp;
        self.generation = self.generation.saturating_add(1);
    }

    /// Returns true if this range is adjacent to another (they could merge).
    ///
    /// Two ranges are adjacent when one's end key equals the other's start key,
    /// and both boundary keys are concrete (not None, which represents infinity).
    #[must_use]
    pub fn is_adjacent_to(&self, other: &RangeDescriptor) -> bool {
        // This range ends where other begins (both must be Some to be a real boundary)
        let this_to_other = match (&self.end_key, &other.start_key) {
            (Some(end), Some(start)) => end == start,
            _ => false,
        };

        // Other range ends where this begins
        let other_to_this = match (&other.end_key, &self.start_key) {
            (Some(end), Some(start)) => end == start,
            _ => false,
        };

        this_to_other || other_to_this
    }

    /// Checks if two adjacent ranges can merge based on combined size threshold.
    #[must_use]
    pub fn can_merge_with(&self, other: &RangeDescriptor, threshold_bytes: u64) -> bool {
        self.is_adjacent_to(other)
            && self.size_bytes.saturating_add(other.size_bytes) < threshold_bytes
    }
}

/// Collection of all range descriptors in the cluster.
///
/// This is the authoritative metadata for the cluster's shard layout.
/// It's propagated via gossip and stored persistently on all nodes.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MetaRange {
    /// All range descriptors indexed by shard ID.
    pub descriptors: BTreeMap<ShardId, RangeDescriptor>,

    /// Version number for the entire meta-range.
    ///
    /// Incremented on any change to any descriptor.
    pub version: u64,

    /// When this meta-range was last updated.
    #[serde(with = "hlc_serde")]
    pub updated_at: HlcTimestamp,
}

impl MetaRange {
    /// Creates an empty meta-range.
    #[must_use]
    pub fn new() -> Self {
        Self { descriptors: BTreeMap::new(), version: 0, updated_at: HlcTimestamp::default() }
    }

    /// Creates a meta-range with initial shards distributed across nodes.
    ///
    /// Shards are assigned to nodes round-robin style.
    #[must_use]
    pub fn with_initial_shards(num_shards: u32, nodes: &[NodeId], replication_factor: u32) -> Self {
        let mut descriptors = BTreeMap::new();
        let rf = replication_factor as usize;

        for shard_id in 0..num_shards {
            // Round-robin replica assignment
            let mut replicas = Vec::with_capacity(rf);
            for i in 0..rf.min(nodes.len()) {
                let node_idx = (shard_id as usize + i) % nodes.len();
                replicas.push(nodes[node_idx]);
            }

            let descriptor = RangeDescriptor::new_full_range(shard_id, replicas);
            descriptors.insert(shard_id, descriptor);
        }

        Self { descriptors, version: 1, updated_at: HlcTimestamp::default() }
    }

    /// Gets a range descriptor by shard ID.
    #[must_use]
    pub fn get(&self, shard_id: ShardId) -> Option<&RangeDescriptor> {
        self.descriptors.get(&shard_id)
    }

    /// Gets a mutable range descriptor by shard ID.
    pub fn get_mut(&mut self, shard_id: ShardId) -> Option<&mut RangeDescriptor> {
        self.descriptors.get_mut(&shard_id)
    }

    /// Inserts or updates a range descriptor.
    pub fn upsert(&mut self, descriptor: RangeDescriptor, timestamp: HlcTimestamp) {
        self.descriptors.insert(descriptor.shard_id, descriptor);
        self.version = self.version.saturating_add(1);
        self.updated_at = timestamp;
    }

    /// Removes a range descriptor.
    pub fn remove(
        &mut self,
        shard_id: ShardId,
        timestamp: HlcTimestamp,
    ) -> Option<RangeDescriptor> {
        let removed = self.descriptors.remove(&shard_id);
        if removed.is_some() {
            self.version = self.version.saturating_add(1);
            self.updated_at = timestamp;
        }
        removed
    }

    /// Returns the total number of shards.
    #[must_use]
    pub fn num_shards(&self) -> usize {
        self.descriptors.len()
    }

    /// Returns all shard IDs.
    #[must_use]
    pub fn shard_ids(&self) -> Vec<ShardId> {
        self.descriptors.keys().copied().collect()
    }

    /// Finds all shards assigned to a specific node.
    #[must_use]
    pub fn shards_for_node(&self, node_id: NodeId) -> Vec<ShardId> {
        self.descriptors
            .iter()
            .filter_map(
                |(&shard_id, desc)| {
                    if desc.replicas.contains(&node_id) {
                        Some(shard_id)
                    } else {
                        None
                    }
                },
            )
            .collect()
    }

    /// Finds all shards where a node is the leader.
    #[must_use]
    pub fn leader_shards_for_node(&self, node_id: NodeId) -> Vec<ShardId> {
        self.descriptors
            .iter()
            .filter_map(
                |(&shard_id, desc)| {
                    if desc.leader() == Some(node_id) {
                        Some(shard_id)
                    } else {
                        None
                    }
                },
            )
            .collect()
    }

    /// Merges another meta-range into this one, keeping newer descriptors.
    ///
    /// Used during gossip to merge remote state.
    pub fn merge(&mut self, other: &MetaRange) {
        for (shard_id, other_desc) in &other.descriptors {
            match self.descriptors.get(shard_id) {
                Some(our_desc) if our_desc.generation >= other_desc.generation => {
                    // Our version is newer or equal, keep ours
                }
                _ => {
                    // Other version is newer, take theirs
                    self.descriptors.insert(*shard_id, other_desc.clone());
                }
            }
        }

        if other.version > self.version {
            self.version = other.version;
            self.updated_at = other.updated_at;
        }
    }
}

impl Default for MetaRange {
    fn default() -> Self {
        Self::new()
    }
}

/// Role of a node for a specific shard.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShardRole {
    /// This node is the leader for the shard.
    Leader,
    /// This node is a follower for the shard.
    Follower,
    /// This node is not a replica for the shard.
    None,
}

impl RangeDescriptor {
    /// Returns this node's role for this shard.
    #[must_use]
    pub fn role_for_node(&self, node_id: NodeId) -> ShardRole {
        if self.leader() == Some(node_id) {
            ShardRole::Leader
        } else if self.replicas.contains(&node_id) {
            ShardRole::Follower
        } else {
            ShardRole::None
        }
    }
}

#[cfg(test)]
#[path = "types_tests.rs"]
mod tests;
