//! Range router for subject-to-shard mapping.
//!
//! This module provides consistent hashing to route subjects to shards
//! using Google's jump hash algorithm for minimal disruption during
//! cluster resizing.

use dashmap::DashMap;
use parking_lot::RwLock;
use std::sync::Arc;
use tracing::instrument;

use crate::membership::NodeId;
use crate::sharding::types::{MetaRange, RangeDescriptor, ShardId};
use crate::{ClusterError, Result};

/// Routes subjects to shards and tracks shard-to-node mappings.
///
/// The router maintains a cached view of the cluster's meta-range and
/// provides efficient subject→shard→nodes lookups.
pub struct RangeRouter {
    /// Cached meta-range (authoritative shard metadata).
    meta_range: RwLock<MetaRange>,

    /// Local node ID (used for preferring local replicas).
    local_node_id: NodeId,

    /// Cached shard-to-replicas mapping for fast lookups.
    replica_cache: DashMap<ShardId, Vec<NodeId>>,
}

impl RangeRouter {
    /// Creates a new range router with the given local node ID.
    pub fn new(local_node_id: NodeId) -> Self {
        Self {
            meta_range: RwLock::new(MetaRange::new()),
            local_node_id,
            replica_cache: DashMap::new(),
        }
    }

    /// Creates a range router with an initial meta-range.
    pub fn with_meta_range(local_node_id: NodeId, meta_range: MetaRange) -> Self {
        let router = Self::new(local_node_id);
        router.update_meta_range(meta_range);
        router
    }

    /// Routes a subject string to its shard ID using jump hash.
    ///
    /// This uses BLAKE3 to hash the subject and Google's jump hash
    /// algorithm for consistent distribution with minimal disruption
    /// when the number of shards changes.
    ///
    /// # Errors
    ///
    /// Returns `ClusterError::Sharding` if no shards are configured.
    #[instrument(skip(self), fields(subject_len = subject.len()))]
    pub fn route_subject(&self, subject: &str) -> Result<ShardId> {
        let hash = blake3::hash(subject.as_bytes());
        let key = u64::from_le_bytes(hash.as_bytes()[0..8].try_into().unwrap_or([0u8; 8]));

        let num_shards = self.num_shards();
        if num_shards == 0 {
            return Err(ClusterError::Sharding("No shards configured".to_string()));
        }

        Ok(jump_hash(key, num_shards))
    }

    /// Routes a raw key (bytes) to its shard ID.
    ///
    /// # Errors
    ///
    /// Returns `ClusterError::Sharding` if no shards are configured.
    pub fn route_key(&self, key: &[u8]) -> Result<ShardId> {
        let hash = blake3::hash(key);
        let hash_u64 = u64::from_le_bytes(hash.as_bytes()[0..8].try_into().unwrap_or([0u8; 8]));

        let num_shards = self.num_shards();
        if num_shards == 0 {
            return Err(ClusterError::Sharding("No shards configured".to_string()));
        }

        Ok(jump_hash(hash_u64, num_shards))
    }

    /// Gets the replicas for a shard, preferring the local node if it's a replica.
    #[instrument(skip(self))]
    pub fn get_replicas(&self, shard_id: ShardId) -> Result<Vec<NodeId>> {
        // Check cache first
        if let Some(replicas) = self.replica_cache.get(&shard_id) {
            return Ok(replicas.clone());
        }

        // Lookup from meta-range
        let meta = self.meta_range.read();
        let descriptor = meta.get(shard_id).ok_or(ClusterError::ShardNotFound(shard_id))?;

        let replicas = descriptor.replicas.clone();

        // Cache the result
        drop(meta);
        self.replica_cache.insert(shard_id, replicas.clone());

        Ok(replicas)
    }

    /// Gets the replicas for a shard, with the local node first if present.
    ///
    /// This is useful for read operations where we prefer local data.
    #[instrument(skip(self))]
    pub fn get_replicas_prefer_local(&self, shard_id: ShardId) -> Result<Vec<NodeId>> {
        let replicas = self.get_replicas(shard_id)?;

        // If local node is a replica, move it to front
        if replicas.contains(&self.local_node_id) {
            let mut reordered = vec![self.local_node_id];
            for node in replicas {
                if node != self.local_node_id {
                    reordered.push(node);
                }
            }
            Ok(reordered)
        } else {
            Ok(replicas)
        }
    }

    /// Gets the leader node for a shard.
    #[instrument(skip(self))]
    pub fn get_leader(&self, shard_id: ShardId) -> Result<NodeId> {
        let meta = self.meta_range.read();
        let descriptor = meta.get(shard_id).ok_or(ClusterError::ShardNotFound(shard_id))?;

        descriptor.leader().ok_or(ClusterError::NoReplicasAvailable(shard_id))
    }

    /// Gets the range descriptor for a shard.
    pub fn get_descriptor(&self, shard_id: ShardId) -> Result<RangeDescriptor> {
        let meta = self.meta_range.read();
        meta.get(shard_id).cloned().ok_or(ClusterError::ShardNotFound(shard_id))
    }

    /// Updates the meta-range and invalidates caches.
    #[instrument(skip(self, meta_range), fields(version = meta_range.version))]
    pub fn update_meta_range(&self, meta_range: MetaRange) {
        // Clear cache before updating
        self.replica_cache.clear();

        let mut current = self.meta_range.write();
        *current = meta_range;
    }

    /// Merges a remote meta-range into the current one.
    #[instrument(skip(self, remote), fields(remote_version = remote.version))]
    pub fn merge_meta_range(&self, remote: &MetaRange) {
        // Clear cache before merging
        self.replica_cache.clear();

        let mut current = self.meta_range.write();
        current.merge(remote);
    }

    /// Returns the current number of shards.
    pub fn num_shards(&self) -> u32 {
        let meta = self.meta_range.read();
        meta.num_shards() as u32
    }

    /// Returns the current meta-range version.
    pub fn version(&self) -> u64 {
        let meta = self.meta_range.read();
        meta.version
    }

    /// Returns a clone of the current meta-range.
    pub fn get_meta_range(&self) -> MetaRange {
        let meta = self.meta_range.read();
        meta.clone()
    }

    /// Returns all shards that this node is a replica for.
    pub fn local_shards(&self) -> Vec<ShardId> {
        let meta = self.meta_range.read();
        meta.shards_for_node(self.local_node_id)
    }

    /// Returns all shards that this node is the leader for.
    pub fn leader_shards(&self) -> Vec<ShardId> {
        let meta = self.meta_range.read();
        meta.leader_shards_for_node(self.local_node_id)
    }

    /// Checks if this node is a replica for the given shard.
    pub fn is_replica_for(&self, shard_id: ShardId) -> bool {
        if let Ok(replicas) = self.get_replicas(shard_id) {
            replicas.contains(&self.local_node_id)
        } else {
            false
        }
    }

    /// Invalidates cached replica entries containing the given node.
    ///
    /// Call this when a node fails or leaves the cluster so that stale
    /// replica lists are evicted from the cache.
    pub fn invalidate_node(&self, node_id: NodeId) {
        self.replica_cache.retain(|_, replicas| !replicas.contains(&node_id));
    }

    /// Checks if this node is the leader for the given shard.
    pub fn is_leader_for(&self, shard_id: ShardId) -> bool {
        if let Ok(leader) = self.get_leader(shard_id) {
            leader == self.local_node_id
        } else {
            false
        }
    }
}

/// Google's jump consistent hash algorithm.
///
/// Maps a key to one of `num_buckets` buckets with:
/// - O(1) time complexity
/// - O(1) space complexity
/// - Minimal disruption when bucket count changes
///
/// Reference: "A Fast, Minimal Memory, Consistent Hash Algorithm"
/// https://arxiv.org/abs/1406.2294
fn jump_hash(key: u64, num_buckets: u32) -> u32 {
    let mut k = key;
    let mut b: i64 = -1;
    let mut j: i64 = 0;

    while j < num_buckets as i64 {
        b = j;
        k = k.wrapping_mul(2862933555777941757).wrapping_add(1);
        j = ((b.wrapping_add(1) as f64)
            * (((1u64 << 31) as f64) / (((k >> 33).wrapping_add(1)) as f64))) as i64;
    }

    b as u32
}

/// Thread-safe wrapper around RangeRouter.
pub type SharedRangeRouter = Arc<RangeRouter>;

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

    fn test_node_id(n: u8) -> NodeId {
        NodeId::from_bytes([n; 16])
    }

    #[test]
    fn test_jump_hash_distribution() {
        // Test that jump hash distributes evenly
        let num_buckets = 10u32;
        let mut bucket_counts = vec![0u64; num_buckets as usize];

        for i in 0..10000u64 {
            let bucket = jump_hash(i, num_buckets);
            bucket_counts[bucket as usize] += 1;
        }

        // Each bucket should have roughly 1000 items (10%)
        // Allow 20% variance
        for count in bucket_counts {
            assert!(count > 800, "Bucket has too few items: {count}");
            assert!(count < 1200, "Bucket has too many items: {count}");
        }
    }

    #[test]
    fn test_jump_hash_consistency() {
        // Same key should always map to same bucket
        let key = 12345u64;
        let bucket1 = jump_hash(key, 10);
        let bucket2 = jump_hash(key, 10);
        assert_eq!(bucket1, bucket2);
    }

    #[test]
    fn test_jump_hash_stability() {
        // Most keys should stay in same bucket when adding a bucket
        let mut unchanged = 0;
        let old_buckets = 10u32;
        let new_buckets = 11u32;

        for i in 0..10000u64 {
            let old_bucket = jump_hash(i, old_buckets);
            let new_bucket = jump_hash(i, new_buckets);
            if old_bucket == new_bucket {
                unchanged += 1;
            }
        }

        // At least 90% should be unchanged (ideally ~91%)
        assert!(unchanged > 9000, "Too many keys moved: {unchanged}/10000 unchanged");
    }

    #[test]
    fn test_route_subject_consistency() {
        let router = RangeRouter::new(test_node_id(1));

        // Initialize with some shards
        let nodes = vec![test_node_id(1), test_node_id(2), test_node_id(3)];
        let meta = MetaRange::with_initial_shards(8, &nodes, 2);
        router.update_meta_range(meta);

        // Same subject should always route to same shard
        let shard1 = router.route_subject("test:subject:123").unwrap();
        let shard2 = router.route_subject("test:subject:123").unwrap();
        assert_eq!(shard1, shard2);
    }

    #[test]
    fn test_get_replicas() {
        let router = RangeRouter::new(test_node_id(1));

        let nodes = vec![test_node_id(1), test_node_id(2), test_node_id(3)];
        let meta = MetaRange::with_initial_shards(4, &nodes, 2);
        router.update_meta_range(meta);

        let replicas = router.get_replicas(0).unwrap();
        assert_eq!(replicas.len(), 2);
    }

    #[test]
    fn test_get_replicas_prefer_local() {
        let local_node = test_node_id(2);
        let router = RangeRouter::new(local_node);

        // Create meta where node 2 is a follower for shard 0
        let nodes = vec![test_node_id(1), test_node_id(2), test_node_id(3)];
        let meta = MetaRange::with_initial_shards(4, &nodes, 2);
        router.update_meta_range(meta);

        // For any shard where local node is a replica, it should be first
        for shard_id in 0..4 {
            let replicas = router.get_replicas(shard_id).unwrap();
            let preferred = router.get_replicas_prefer_local(shard_id).unwrap();

            if replicas.contains(&local_node) {
                assert_eq!(preferred[0], local_node);
            }
        }
    }

    #[test]
    fn test_local_shards() {
        let local_node = test_node_id(1);
        let router = RangeRouter::new(local_node);

        let nodes = vec![test_node_id(1), test_node_id(2), test_node_id(3)];
        let meta = MetaRange::with_initial_shards(6, &nodes, 2);
        router.update_meta_range(meta);

        let local_shards = router.local_shards();
        // With round-robin and RF=2, node 1 should be replica for multiple shards
        assert!(!local_shards.is_empty());
    }

    #[test]
    fn test_shard_not_found() {
        let router = RangeRouter::new(test_node_id(1));
        router.update_meta_range(MetaRange::new());

        let result = router.get_replicas(999);
        assert!(matches!(result, Err(ClusterError::ShardNotFound(999))));
    }

    #[test]
    fn test_merge_meta_range() {
        let router = RangeRouter::new(test_node_id(1));

        let nodes = vec![test_node_id(1), test_node_id(2)];
        let meta1 = MetaRange::with_initial_shards(2, &nodes, 2);
        router.update_meta_range(meta1);

        let initial_version = router.version();

        // Create updated meta with higher version
        let mut meta2 = router.get_meta_range();
        if let Some(desc) = meta2.get_mut(0) {
            desc.size_bytes = 5000;
            desc.generation = 100;
        }
        meta2.version = initial_version + 10;

        router.merge_meta_range(&meta2);

        // Version should be updated
        assert!(router.version() > initial_version);

        // Descriptor should have new data
        let desc = router.get_descriptor(0).unwrap();
        assert_eq!(desc.size_bytes, 5000);
    }

    #[test]
    fn test_route_subject_no_shards() {
        let router = RangeRouter::new(test_node_id(1));
        // Empty meta-range: no shards configured
        router.update_meta_range(MetaRange::new());

        let result = router.route_subject("test:subject");
        assert!(result.is_err());
    }

    #[test]
    fn test_invalidate_node() {
        let router = RangeRouter::new(test_node_id(1));

        let nodes = vec![test_node_id(1), test_node_id(2), test_node_id(3)];
        let meta = MetaRange::with_initial_shards(4, &nodes, 2);
        router.update_meta_range(meta);

        // Populate cache
        let _ = router.get_replicas(0);
        let _ = router.get_replicas(1);
        assert!(!router.replica_cache.is_empty());

        // Invalidate node 2 - should evict any cached entries containing it
        router.invalidate_node(test_node_id(2));

        // Cache entries containing node 2 should be gone; re-fetching works
        let replicas = router.get_replicas(0).unwrap();
        assert!(!replicas.is_empty());
    }
}
