lcpfs 2026.1.102

//! CRUSH - Controlled Replication Under Scalable Hashing
//!
//! This module implements the CRUSH algorithm for deterministic, pseudo-random
//! data placement in distributed storage systems. CRUSH allows clients to
//! calculate object locations directly without querying a central directory.
//!
//! # Algorithm Overview
//!
//! CRUSH uses a hierarchical cluster map and placement rules to select
//! storage targets (OSDs) for each object. The selection is:
//!
//! - **Deterministic**: Same input always produces same output
//! - **Pseudo-random**: Appears random but is reproducible
//! - **Weighted**: Respects device capacity differences
//! - **Failure-domain aware**: Spreads replicas across failure boundaries
//!
//! # Hierarchy
//!
//! ```text
//! root (datacenter)
//! ├── rack-1
//! │   ├── host-1
//! │   │   ├── osd.0
//! │   │   └── osd.1
//! │   └── host-2
//! │       ├── osd.2
//! │       └── osd.3
//! └── rack-2
//!     ├── host-3
//!     │   ├── osd.4
//!     │   └── osd.5
//!     └── host-4
//!         ├── osd.6
//!         └── osd.7
//! ```
//!
//! # References
//!
//! - Weil, S. A., et al. "CRUSH: Controlled, Scalable, Decentralized Placement
//!   of Replicated Data." SC '06 Proceedings.

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use core::fmt;

// ============================================================================
// CRUSH Hash Functions
// ============================================================================

/// Robert Jenkins' 32-bit integer hash function.
/// Provides excellent avalanche properties for CRUSH selection.
#[inline]
pub fn crush_hash32(mut a: u32) -> u32 {
    a = a.wrapping_add(!(a << 15));
    a ^= a >> 10;
    a = a.wrapping_add(a << 3);
    a ^= a >> 6;
    a = a.wrapping_add(!(a << 11));
    a ^= a >> 16;
    a
}

/// 64-bit hash by combining two 32-bit hashes.
#[inline]
pub fn crush_hash64(x: u64) -> u64 {
    let low = crush_hash32(x as u32);
    let high = crush_hash32((x >> 32) as u32);
    ((high as u64) << 32) | (low as u64)
}

/// Hash with multiple inputs for CRUSH selection.
/// Used to generate different selections for different replicas.
#[inline]
pub fn crush_hash(x: u64, r: u64) -> u64 {
    let mut h = x;
    h = h.wrapping_mul(0x87c37b91114253d5);
    h = h.wrapping_add(r);
    h ^= h >> 33;
    h = h.wrapping_mul(0x4cf5ad432745937f);
    h ^= h >> 29;
    h = h.wrapping_mul(0x94d049bb133111eb);
    h ^= h >> 32;
    h
}

/// Hash with three inputs (x, bucket_id, replica).
#[inline]
pub fn crush_hash3(x: u64, b: i64, r: u64) -> u64 {
    crush_hash(crush_hash(x, b as u64), r)
}

// ============================================================================
// Bucket Types and Structures
// ============================================================================

/// Type of bucket in the CRUSH hierarchy.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BucketType {
    /// Individual OSD (leaf node)
    Osd,
    /// Physical host/server
    Host,
    /// Network rack
    Rack,
    /// Room or pod
    Room,
    /// Data center
    Datacenter,
    /// Geographic region
    Region,
    /// Root of the hierarchy
    Root,
    /// Custom type
    Custom(u8),
}

impl BucketType {
    /// Get the numeric ID for this bucket type.
    /// Higher values = higher in hierarchy.
    pub fn type_id(&self) -> i32 {
        match self {
            BucketType::Osd => 0,
            BucketType::Host => 1,
            BucketType::Rack => 2,
            BucketType::Room => 3,
            BucketType::Datacenter => 4,
            BucketType::Region => 5,
            BucketType::Root => 6,
            BucketType::Custom(id) => 100 + (*id as i32),
        }
    }

    /// Parse from type ID.
    pub fn from_type_id(id: i32) -> Self {
        match id {
            0 => BucketType::Osd,
            1 => BucketType::Host,
            2 => BucketType::Rack,
            3 => BucketType::Room,
            4 => BucketType::Datacenter,
            5 => BucketType::Region,
            6 => BucketType::Root,
            x if x >= 100 => BucketType::Custom((x - 100) as u8),
            _ => BucketType::Custom(0),
        }
    }
}

impl fmt::Display for BucketType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BucketType::Osd => write!(f, "osd"),
            BucketType::Host => write!(f, "host"),
            BucketType::Rack => write!(f, "rack"),
            BucketType::Room => write!(f, "room"),
            BucketType::Datacenter => write!(f, "datacenter"),
            BucketType::Region => write!(f, "region"),
            BucketType::Root => write!(f, "root"),
            BucketType::Custom(id) => write!(f, "custom-{}", id),
        }
    }
}

/// Selection algorithm for a bucket.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum BucketAlgorithm {
    /// Uniform random selection (unweighted).
    Uniform,
    /// List bucket (ordered).
    List,
    /// Tree bucket (balanced tree).
    Tree,
    /// Straw bucket (original).
    Straw,
    /// Straw2 bucket (improved, default).
    #[default]
    Straw2,
}

/// A bucket in the CRUSH hierarchy.
/// Can contain other buckets or OSDs.
#[derive(Debug, Clone)]
pub struct CrushBucket {
    /// Unique bucket ID (negative for buckets, positive for OSDs).
    pub id: i64,
    /// Human-readable name.
    pub name: String,
    /// Type of this bucket.
    pub bucket_type: BucketType,
    /// Selection algorithm.
    pub algorithm: BucketAlgorithm,
    /// Hash seed for this bucket.
    pub hash_seed: u32,
    /// Weight of this bucket (sum of children for non-leaf).
    pub weight: u32,
    /// Child bucket/OSD IDs.
    pub children: Vec<i64>,
    /// Child weights (parallel to children).
    pub child_weights: Vec<u32>,
}

impl CrushBucket {
    /// Create a new bucket.
    pub fn new(id: i64, name: &str, bucket_type: BucketType) -> Self {
        Self {
            id,
            name: name.to_string(),
            bucket_type,
            algorithm: BucketAlgorithm::Straw2,
            hash_seed: id as u32,
            weight: 0,
            children: Vec::new(),
            child_weights: Vec::new(),
        }
    }

    /// Add a child with weight.
    pub fn add_child(&mut self, child_id: i64, weight: u32) {
        self.children.push(child_id);
        self.child_weights.push(weight);
        self.weight += weight;
    }

    /// Remove a child.
    pub fn remove_child(&mut self, child_id: i64) -> bool {
        if let Some(pos) = self.children.iter().position(|&id| id == child_id) {
            self.weight -= self.child_weights[pos];
            self.children.remove(pos);
            self.child_weights.remove(pos);
            true
        } else {
            false
        }
    }

    /// Check if this is a leaf (OSD).
    pub fn is_leaf(&self) -> bool {
        self.bucket_type == BucketType::Osd
    }

    /// Get the number of children.
    pub fn size(&self) -> usize {
        self.children.len()
    }
}

// ============================================================================
// CRUSH Rules
// ============================================================================

/// A step in a CRUSH rule.
#[derive(Debug, Clone)]
pub enum CrushStep {
    /// Take a specific bucket as starting point.
    Take(i64),
    /// Choose N items of a given type.
    Choose {
        /// Number of items to select.
        count: usize,
        /// Type of item to select.
        bucket_type: BucketType,
    },
    /// Choose N items, allowing selection from same parent.
    ChooseLeaf {
        /// Number of items to select.
        count: usize,
        /// Type of failure domain.
        bucket_type: BucketType,
    },
    /// Choose the first N available items.
    ChooseFirstN {
        /// Number of items to select.
        count: usize,
        /// Type of item to select.
        bucket_type: BucketType,
    },
    /// Emit the current selection as results.
    Emit,
    /// Set the choose mode for subsequent steps.
    SetChooseMode(ChooseMode),
}

/// Mode for choosing items.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ChooseMode {
    /// First N algorithm.
    FirstN,
    /// Indep algorithm (more stable on changes).
    #[default]
    Indep,
}

/// A CRUSH placement rule.
#[derive(Debug, Clone)]
pub struct CrushRule {
    /// Rule ID.
    pub id: u32,
    /// Rule name.
    pub name: String,
    /// Minimum replication factor.
    pub min_size: usize,
    /// Maximum replication factor.
    pub max_size: usize,
    /// Steps in this rule.
    pub steps: Vec<CrushStep>,
}

impl CrushRule {
    /// Create a new rule.
    pub fn new(id: u32, name: &str) -> Self {
        Self {
            id,
            name: name.to_string(),
            min_size: 1,
            max_size: 10,
            steps: Vec::new(),
        }
    }

    /// Add a step to the rule.
    pub fn add_step(&mut self, step: CrushStep) {
        self.steps.push(step);
    }

    /// Create a simple replicated rule.
    pub fn replicated(id: u32, name: &str, root_id: i64, failure_domain: BucketType) -> Self {
        let mut rule = Self::new(id, name);
        rule.steps = vec![
            CrushStep::Take(root_id),
            CrushStep::ChooseLeaf {
                count: 0, // 0 means use pool's replication factor
                bucket_type: failure_domain,
            },
            CrushStep::Emit,
        ];
        rule
    }

    /// Create an erasure-coded rule.
    pub fn erasure(id: u32, name: &str, root_id: i64, failure_domain: BucketType) -> Self {
        let mut rule = Self::new(id, name);
        rule.steps = vec![
            CrushStep::Take(root_id),
            CrushStep::Choose {
                count: 0,
                bucket_type: failure_domain,
            },
            CrushStep::ChooseLeaf {
                count: 1,
                bucket_type: BucketType::Osd,
            },
            CrushStep::Emit,
        ];
        rule
    }
}

// ============================================================================
// CRUSH Map
// ============================================================================

/// The complete CRUSH map for a cluster.
#[derive(Debug, Clone)]
pub struct CrushMap {
    /// All buckets indexed by ID.
    buckets: BTreeMap<i64, CrushBucket>,
    /// All rules indexed by ID.
    rules: BTreeMap<u32, CrushRule>,
    /// Rules indexed by name.
    rules_by_name: BTreeMap<String, u32>,
    /// OSD weights (separate from bucket structure).
    weights: BTreeMap<u64, f64>,
    /// Root bucket ID.
    root_id: i64,
    /// Tunables.
    tunables: CrushTunables,
}

/// CRUSH tunables for algorithm behavior.
#[derive(Debug, Clone)]
pub struct CrushTunables {
    /// Use Straw2 algorithm.
    pub straw2: bool,
    /// Choose local retry count.
    pub choose_local_tries: u32,
    /// Choose local fallback count.
    pub choose_local_fallback_tries: u32,
    /// Total choose tries.
    pub choose_total_tries: u32,
    /// Allow retry descent.
    pub chooseleaf_descend_once: bool,
    /// Vary replica selection.
    pub chooseleaf_vary_r: bool,
    /// Stable selection on OSD removal.
    pub chooseleaf_stable: bool,
}

impl Default for CrushTunables {
    fn default() -> Self {
        Self {
            straw2: true,
            choose_local_tries: 2,
            choose_local_fallback_tries: 5,
            choose_total_tries: 50,
            chooseleaf_descend_once: true,
            chooseleaf_vary_r: true,
            chooseleaf_stable: true,
        }
    }
}

/// Error during CRUSH operation.
#[derive(Debug, Clone)]
pub enum CrushError {
    /// Bucket not found.
    BucketNotFound(i64),
    /// Rule not found.
    RuleNotFound(String),
    /// Not enough OSDs for replication.
    InsufficientOsds {
        /// Required number of OSDs.
        required: usize,
        /// Available number of OSDs.
        available: usize,
    },
    /// Invalid rule configuration.
    InvalidRule(String),
    /// Weight is zero.
    ZeroWeight,
}

impl fmt::Display for CrushError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CrushError::BucketNotFound(id) => write!(f, "Bucket {} not found", id),
            CrushError::RuleNotFound(name) => write!(f, "Rule '{}' not found", name),
            CrushError::InsufficientOsds {
                required,
                available,
            } => {
                write!(f, "Need {} OSDs but only {} available", required, available)
            }
            CrushError::InvalidRule(msg) => write!(f, "Invalid rule: {}", msg),
            CrushError::ZeroWeight => write!(f, "Cannot select from zero-weight bucket"),
        }
    }
}

impl CrushMap {
    /// Create a new empty CRUSH map.
    pub fn new() -> Self {
        Self {
            buckets: BTreeMap::new(),
            rules: BTreeMap::new(),
            rules_by_name: BTreeMap::new(),
            weights: BTreeMap::new(),
            root_id: -1,
            tunables: CrushTunables::default(),
        }
    }

    /// Create a simple CRUSH map with single host.
    pub fn simple(num_osds: usize) -> Self {
        let mut map = Self::new();

        // Create root bucket
        let root_id = -1;
        let mut root = CrushBucket::new(root_id, "root", BucketType::Root);

        // Create host bucket
        let host_id = -2;
        let mut host = CrushBucket::new(host_id, "host0", BucketType::Host);

        // Add OSDs to host
        for i in 0..num_osds {
            let osd_id = i as i64;
            let weight = 0x10000; // 1.0 in fixed-point
            host.add_child(osd_id, weight);
            map.weights.insert(i as u64, 1.0);
        }

        // Add host to root
        root.add_child(host_id, host.weight);

        map.buckets.insert(host_id, host);
        map.buckets.insert(root_id, root);
        map.root_id = root_id;

        // Create default replicated rule
        let rule = CrushRule::replicated(0, "replicated_rule", root_id, BucketType::Host);
        map.add_rule(rule);

        map
    }

    /// Add a bucket to the map.
    pub fn add_bucket(&mut self, bucket: CrushBucket) {
        self.buckets.insert(bucket.id, bucket);
    }

    /// Get a bucket by ID.
    pub fn get_bucket(&self, id: i64) -> Option<&CrushBucket> {
        self.buckets.get(&id)
    }

    /// Get a mutable bucket by ID.
    pub fn get_bucket_mut(&mut self, id: i64) -> Option<&mut CrushBucket> {
        self.buckets.get_mut(&id)
    }

    /// Add a rule to the map.
    pub fn add_rule(&mut self, rule: CrushRule) {
        self.rules_by_name.insert(rule.name.clone(), rule.id);
        self.rules.insert(rule.id, rule);
    }

    /// Get a rule by name.
    pub fn get_rule(&self, name: &str) -> Option<&CrushRule> {
        self.rules_by_name
            .get(name)
            .and_then(|id| self.rules.get(id))
    }

    /// Set OSD weight.
    pub fn set_weight(&mut self, osd_id: u64, weight: f64) {
        self.weights.insert(osd_id, weight);
    }

    /// Get OSD weight.
    pub fn get_weight(&self, osd_id: u64) -> f64 {
        *self.weights.get(&osd_id).unwrap_or(&1.0)
    }

    /// Add an OSD to the CRUSH map under a host bucket.
    /// If the host doesn't exist, creates a new host bucket.
    /// If no root exists, creates a root bucket.
    pub fn add_osd(&mut self, osd_id: u64, host_name: &str, weight: f64) {
        let weight_fixed = (weight * 65536.0) as u32; // Convert to fixed-point

        // Ensure root exists
        if !self.buckets.contains_key(&self.root_id) {
            let root = CrushBucket::new(-1, "root", BucketType::Root);
            self.buckets.insert(-1, root);
            self.root_id = -1;

            // Create default rule if none exists
            if self.rules.is_empty() {
                let rule = CrushRule::replicated(0, "replicated_rule", -1, BucketType::Host);
                self.add_rule(rule);
            }
        }

        // Find or create host bucket
        let host_id = self.find_or_create_host(host_name);

        // Add OSD to host bucket and get new host weight
        let new_host_weight = {
            if let Some(host) = self.buckets.get_mut(&host_id) {
                // Check if OSD already exists in host
                if !host.children.contains(&(osd_id as i64)) {
                    host.add_child(osd_id as i64, weight_fixed);
                }
                Some(host.weight)
            } else {
                None
            }
        };

        // Update root's weight for this host
        if let Some(host_weight) = new_host_weight {
            if let Some(root) = self.buckets.get_mut(&self.root_id) {
                if let Some(pos) = root.children.iter().position(|&id| id == host_id) {
                    root.weight -= root.child_weights[pos];
                    root.child_weights[pos] = host_weight;
                    root.weight += host_weight;
                }
            }
        }

        // Set weight in weight map
        self.weights.insert(osd_id, weight);
    }

    /// Find or create a host bucket, returning its ID.
    fn find_or_create_host(&mut self, host_name: &str) -> i64 {
        // Look for existing host
        for (&id, bucket) in &self.buckets {
            if bucket.bucket_type == BucketType::Host && bucket.name == host_name {
                return id;
            }
        }

        // Create new host bucket with a negative ID
        let host_id = self
            .buckets
            .keys()
            .filter(|&&k| k < 0)
            .min()
            .map(|&m| m - 1)
            .unwrap_or(-2);

        let host = CrushBucket::new(host_id, host_name, BucketType::Host);
        self.buckets.insert(host_id, host);

        // Add host to root
        if let Some(root) = self.buckets.get_mut(&self.root_id) {
            root.add_child(host_id, 0); // Weight will be updated when OSDs are added
        }

        host_id
    }

    /// Remove an OSD from the CRUSH map.
    pub fn remove_osd(&mut self, osd_id: u64) {
        // Find and remove from parent host
        for bucket in self.buckets.values_mut() {
            if bucket.bucket_type == BucketType::Host && bucket.remove_child(osd_id as i64) {
                break;
            }
        }

        self.weights.remove(&osd_id);
    }

    /// Set the root bucket ID.
    pub fn set_root(&mut self, root_id: i64) {
        self.root_id = root_id;
    }

    /// Get tunables.
    pub fn tunables(&self) -> &CrushTunables {
        &self.tunables
    }

    /// Set tunables.
    pub fn set_tunables(&mut self, tunables: CrushTunables) {
        self.tunables = tunables;
    }

    // ========================================================================
    // Core Selection Algorithm
    // ========================================================================

    /// Select N OSDs for placement using the named rule.
    pub fn select(&self, rule_name: &str, pgid: u64, count: usize) -> Result<Vec<u64>, CrushError> {
        let rule = self
            .get_rule(rule_name)
            .ok_or_else(|| CrushError::RuleNotFound(rule_name.to_string()))?;

        let mut out = Vec::new();
        let mut working = Vec::new();

        for step in &rule.steps {
            match step {
                CrushStep::Take(bucket_id) => {
                    working.clear();
                    working.push(*bucket_id);
                }

                CrushStep::Choose {
                    count: n,
                    bucket_type,
                }
                | CrushStep::ChooseLeaf {
                    count: n,
                    bucket_type,
                } => {
                    let num_to_select = if *n == 0 { count } else { *n };
                    let mut new_working = Vec::new();

                    for &start_bucket in &working {
                        let selected = self.do_choose(
                            start_bucket,
                            pgid,
                            num_to_select,
                            *bucket_type,
                            matches!(step, CrushStep::ChooseLeaf { .. }),
                            &out,
                        )?;
                        new_working.extend(selected);
                    }

                    working = new_working;
                }

                CrushStep::ChooseFirstN {
                    count: n,
                    bucket_type,
                } => {
                    let num_to_select = if *n == 0 { count } else { *n };
                    let mut new_working = Vec::new();

                    for &start_bucket in &working {
                        let selected =
                            self.do_choose_firstn(start_bucket, pgid, num_to_select, *bucket_type)?;
                        new_working.extend(selected);
                    }

                    working = new_working;
                }

                CrushStep::Emit => {
                    // Convert working bucket IDs to OSD IDs
                    for &id in &working {
                        if id >= 0 {
                            out.push(id as u64);
                        }
                    }
                    working.clear();
                }

                CrushStep::SetChooseMode(_) => {
                    // Handled in do_choose
                }
            }
        }

        // Deduplicate results
        let mut seen = alloc::collections::BTreeSet::new();
        out.retain(|&x| seen.insert(x));

        Ok(out)
    }

    /// Internal choose implementation.
    fn do_choose(
        &self,
        start: i64,
        pgid: u64,
        count: usize,
        target_type: BucketType,
        leaf: bool,
        already_selected: &[u64],
    ) -> Result<Vec<i64>, CrushError> {
        let bucket = self
            .buckets
            .get(&start)
            .ok_or(CrushError::BucketNotFound(start))?;

        let mut out = Vec::with_capacity(count);
        let mut collisions = 0;
        let max_tries = self.tunables.choose_total_tries as usize;

        for replica in 0..count {
            let mut r = replica;
            let mut retry = 0;

            loop {
                if retry >= max_tries {
                    break;
                }

                let item = self.select_from_bucket(bucket, pgid, r as u64)?;

                // Check if we need to descend further
                let final_item = if item >= 0 {
                    // It's an OSD
                    item
                } else if let Some(child_bucket) = self.buckets.get(&item) {
                    if child_bucket.bucket_type == target_type {
                        if leaf {
                            // ChooseLeaf: we found the target type, now find an OSD under it
                            self.find_leaf_osd(item, pgid, r as u64)?
                        } else {
                            // Choose: return the bucket itself
                            item
                        }
                    } else {
                        // Descend into this bucket to find target type
                        self.descend_to_type(item, pgid, target_type, leaf, r as u64)?
                    }
                } else {
                    return Err(CrushError::BucketNotFound(item));
                };

                // Check for collision with existing selection
                let osd_id = if final_item >= 0 {
                    final_item as u64
                } else {
                    u64::MAX
                };
                let is_collision = out.contains(&final_item) || already_selected.contains(&osd_id);

                if is_collision {
                    collisions += 1;
                    r += count; // Try different replica slot
                    retry += 1;
                    continue;
                }

                // Check OSD weight (skip if weight is 0)
                if final_item >= 0 {
                    let weight = self.get_weight(final_item as u64);
                    if weight <= 0.0 {
                        r += count;
                        retry += 1;
                        continue;
                    }
                }

                out.push(final_item);
                break;
            }
        }

        Ok(out)
    }

    /// Descend through the hierarchy to find a target type.
    fn descend_to_type(
        &self,
        start: i64,
        pgid: u64,
        target_type: BucketType,
        leaf: bool,
        replica: u64,
    ) -> Result<i64, CrushError> {
        let mut current = start;
        let mut depth = 0;
        const MAX_DEPTH: usize = 20;

        while depth < MAX_DEPTH {
            let bucket = self
                .buckets
                .get(&current)
                .ok_or(CrushError::BucketNotFound(current))?;

            // Found target type
            if bucket.bucket_type == target_type {
                if leaf {
                    // Need to find an OSD under this
                    return self.find_leaf_osd(current, pgid, replica);
                }
                return Ok(current);
            }

            // If we're at an OSD, we're done
            if bucket.bucket_type == BucketType::Osd || bucket.children.is_empty() {
                return Ok(current);
            }

            // Select a child and descend
            let child = self.select_from_bucket(bucket, pgid, replica)?;
            current = child;
            depth += 1;
        }

        Ok(current)
    }

    /// Find a leaf OSD under a bucket.
    fn find_leaf_osd(&self, bucket_id: i64, pgid: u64, replica: u64) -> Result<i64, CrushError> {
        let mut current = bucket_id;
        let mut depth = 0;
        const MAX_DEPTH: usize = 20;

        while depth < MAX_DEPTH {
            if current >= 0 {
                // It's an OSD
                return Ok(current);
            }

            let bucket = self
                .buckets
                .get(&current)
                .ok_or(CrushError::BucketNotFound(current))?;

            if bucket.children.is_empty() {
                return Err(CrushError::InsufficientOsds {
                    required: 1,
                    available: 0,
                });
            }

            current = self.select_from_bucket(bucket, pgid, replica + depth as u64)?;
            depth += 1;
        }

        Ok(current)
    }

    /// Choose using firstn algorithm.
    fn do_choose_firstn(
        &self,
        start: i64,
        pgid: u64,
        count: usize,
        target_type: BucketType,
    ) -> Result<Vec<i64>, CrushError> {
        let bucket = self
            .buckets
            .get(&start)
            .ok_or(CrushError::BucketNotFound(start))?;

        let mut out = Vec::with_capacity(count);

        for (i, &child) in bucket.children.iter().enumerate() {
            if out.len() >= count {
                break;
            }

            if child >= 0 {
                // OSD
                if target_type == BucketType::Osd {
                    out.push(child);
                }
            } else if let Some(child_bucket) = self.buckets.get(&child) {
                if child_bucket.bucket_type == target_type {
                    out.push(child);
                } else {
                    // Descend
                    let descended =
                        self.descend_to_type(child, pgid, target_type, false, i as u64)?;
                    out.push(descended);
                }
            }
        }

        Ok(out)
    }

    // ========================================================================
    // Bucket Selection Algorithms
    // ========================================================================

    /// Select an item from a bucket.
    fn select_from_bucket(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        if bucket.children.is_empty() {
            return Err(CrushError::InsufficientOsds {
                required: 1,
                available: 0,
            });
        }

        match bucket.algorithm {
            BucketAlgorithm::Uniform => self.select_uniform(bucket, x, r),
            BucketAlgorithm::List => self.select_list(bucket, x, r),
            BucketAlgorithm::Tree => self.select_tree(bucket, x, r),
            BucketAlgorithm::Straw => self.select_straw(bucket, x, r),
            BucketAlgorithm::Straw2 => self.select_straw2(bucket, x, r),
        }
    }

    /// Uniform selection (for unweighted buckets).
    fn select_uniform(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        let hash = crush_hash3(x, bucket.id, r);
        let index = (hash as usize) % bucket.children.len();
        Ok(bucket.children[index])
    }

    /// List bucket selection.
    fn select_list(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        let mut remaining = bucket.weight;

        for (i, (&child, &weight)) in bucket
            .children
            .iter()
            .zip(bucket.child_weights.iter())
            .enumerate()
        {
            let hash = crush_hash3(x, bucket.id, r.wrapping_add(i as u64));
            let w = hash % (remaining as u64 + 1);

            if w < weight as u64 {
                return Ok(child);
            }

            remaining -= weight;
        }

        // Fallback to last item
        Ok(*bucket.children.last().unwrap())
    }

    /// Tree bucket selection.
    fn select_tree(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        // Simplified tree selection (full implementation would use a balanced tree)
        self.select_straw2(bucket, x, r)
    }

    /// Straw bucket selection (original algorithm).
    fn select_straw(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        let mut high_draw = 0u64;
        let mut high_item = bucket.children[0];

        for (&child, &weight) in bucket.children.iter().zip(bucket.child_weights.iter()) {
            let hash = crush_hash3(x, child, r);
            let draw = hash.wrapping_mul(weight as u64);

            if draw > high_draw {
                high_draw = draw;
                high_item = child;
            }
        }

        Ok(high_item)
    }

    /// Straw2 bucket selection (improved algorithm).
    /// This is the default and recommended algorithm.
    fn select_straw2(&self, bucket: &CrushBucket, x: u64, r: u64) -> Result<i64, CrushError> {
        let mut high_draw = i64::MIN;
        let mut high_item = bucket.children[0];

        for (&child, &weight) in bucket.children.iter().zip(bucket.child_weights.iter()) {
            if weight == 0 {
                continue;
            }

            let hash = crush_hash3(x, child, r);

            // Straw2 formula: draw = ln(hash/max) / weight
            // We use a simplified version that maintains ordering
            let u = (hash as f64) / (u64::MAX as f64);
            let ln_u = if u > 0.0 { libm::log(u) } else { -1e10 };
            let draw = (ln_u / (weight as f64) * 1e9) as i64;

            if draw > high_draw {
                high_draw = draw;
                high_item = child;
            }
        }

        Ok(high_item)
    }

    // ========================================================================
    // Utility Methods
    // ========================================================================

    /// Get all OSD IDs in the cluster.
    pub fn get_all_osds(&self) -> Vec<u64> {
        let mut osds = Vec::new();
        for (&id, bucket) in &self.buckets {
            if bucket.bucket_type == BucketType::Osd && id >= 0 {
                osds.push(id as u64);
            }
        }
        // Also check children that are OSDs
        for bucket in self.buckets.values() {
            for &child in &bucket.children {
                if child >= 0 && !osds.contains(&(child as u64)) {
                    osds.push(child as u64);
                }
            }
        }
        osds.sort_unstable();
        osds.dedup();
        osds
    }

    /// Get the total weight of all OSDs.
    pub fn total_weight(&self) -> f64 {
        self.weights.values().sum()
    }

    /// Count OSDs in a bucket (recursively).
    pub fn count_osds_in_bucket(&self, bucket_id: i64) -> usize {
        if bucket_id >= 0 {
            return 1; // It's an OSD
        }

        let bucket = match self.buckets.get(&bucket_id) {
            Some(b) => b,
            None => return 0,
        };

        let mut count = 0;
        for &child in &bucket.children {
            count += self.count_osds_in_bucket(child);
        }
        count
    }

    /// Get the failure domain for an OSD.
    pub fn get_failure_domain(&self, osd_id: u64, domain_type: BucketType) -> Option<i64> {
        // Find which bucket of the given type contains this OSD
        for (bucket_id, bucket) in &self.buckets {
            if bucket.bucket_type == domain_type && self.bucket_contains_osd(*bucket_id, osd_id) {
                return Some(*bucket_id);
            }
        }
        None
    }

    /// Check if a bucket (recursively) contains an OSD.
    fn bucket_contains_osd(&self, bucket_id: i64, osd_id: u64) -> bool {
        if bucket_id >= 0 {
            return bucket_id as u64 == osd_id;
        }

        let bucket = match self.buckets.get(&bucket_id) {
            Some(b) => b,
            None => return false,
        };

        for &child in &bucket.children {
            if self.bucket_contains_osd(child, osd_id) {
                return true;
            }
        }

        false
    }
}

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

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_crush_hash_deterministic() {
        // Hash should be deterministic
        let h1 = crush_hash(12345, 0);
        let h2 = crush_hash(12345, 0);
        assert_eq!(h1, h2);

        // Different inputs should give different outputs
        let h3 = crush_hash(12345, 1);
        assert_ne!(h1, h3);
    }

    #[test]
    fn test_crush_hash_distribution() {
        // Check hash distribution across buckets
        let mut buckets = [0u32; 10];
        for i in 0..10000u64 {
            let h = crush_hash(i, 0);
            let bucket = (h % 10) as usize;
            buckets[bucket] += 1;
        }

        // Each bucket should have roughly 1000 items (10% of 10000)
        for &count in &buckets {
            assert!(
                count > 800 && count < 1200,
                "Bucket count {} is not well distributed",
                count
            );
        }
    }

    #[test]
    fn test_simple_crush_map() {
        let map = CrushMap::simple(4);

        assert_eq!(map.get_all_osds().len(), 4);
        assert!(map.get_rule("replicated_rule").is_some());
    }

    #[test]
    fn test_crush_select_deterministic() {
        let map = CrushMap::simple(8);

        let osds1 = map.select("replicated_rule", 100, 3).unwrap();
        let osds2 = map.select("replicated_rule", 100, 3).unwrap();

        assert_eq!(osds1, osds2);
        assert_eq!(osds1.len(), 3);
    }

    #[test]
    fn test_crush_select_unique() {
        let map = CrushMap::simple(8);

        let osds = map.select("replicated_rule", 100, 3).unwrap();

        // All OSDs should be unique
        let mut seen = alloc::collections::BTreeSet::new();
        for osd in &osds {
            assert!(seen.insert(*osd), "Duplicate OSD {} in selection", osd);
        }
    }

    #[test]
    fn test_crush_select_distribution() {
        let map = CrushMap::simple(8);
        let mut osd_counts = [0u32; 8];

        // Select for many PGs
        for pgid in 0..1000u64 {
            let osds = map.select("replicated_rule", pgid, 3).unwrap();
            for osd in osds {
                osd_counts[osd as usize] += 1;
            }
        }

        // Each OSD should get roughly equal share (375 = 3000/8)
        let expected = 375;
        for &count in &osd_counts {
            let deviation = (count as i32 - expected).unsigned_abs();
            assert!(
                deviation < 100,
                "OSD count {} deviates too much from {}",
                count,
                expected
            );
        }
    }

    #[test]
    fn test_crush_select_different_pgs() {
        let map = CrushMap::simple(8);

        let osds1 = map.select("replicated_rule", 1, 3).unwrap();
        let osds2 = map.select("replicated_rule", 2, 3).unwrap();

        // Different PGs should likely get different selections
        // (Not guaranteed but very likely with good distribution)
        let same = osds1 == osds2;
        // This is probabilistic, so we just note it
        assert!(osds1.len() == 3);
        assert!(osds2.len() == 3);
        let _ = same; // Acknowledge we checked it
    }

    #[test]
    fn test_bucket_type_ordering() {
        assert!(BucketType::Root.type_id() > BucketType::Datacenter.type_id());
        assert!(BucketType::Datacenter.type_id() > BucketType::Rack.type_id());
        assert!(BucketType::Rack.type_id() > BucketType::Host.type_id());
        assert!(BucketType::Host.type_id() > BucketType::Osd.type_id());
    }

    #[test]
    fn test_bucket_operations() {
        let mut bucket = CrushBucket::new(-1, "test", BucketType::Host);

        bucket.add_child(0, 100);
        bucket.add_child(1, 100);
        bucket.add_child(2, 100);

        assert_eq!(bucket.size(), 3);
        assert_eq!(bucket.weight, 300);

        bucket.remove_child(1);

        assert_eq!(bucket.size(), 2);
        assert_eq!(bucket.weight, 200);
    }

    #[test]
    fn test_crush_rule_creation() {
        let rule = CrushRule::replicated(0, "test", -1, BucketType::Host);

        assert_eq!(rule.steps.len(), 3);
        assert!(matches!(rule.steps[0], CrushStep::Take(-1)));
        assert!(matches!(rule.steps[2], CrushStep::Emit));
    }

    #[test]
    fn test_insufficient_osds() {
        let map = CrushMap::simple(2);

        // Requesting more replicas than OSDs
        let result = map.select("replicated_rule", 1, 5);

        // Should return what's available (2 OSDs)
        let osds = result.unwrap();
        assert!(osds.len() <= 2);
    }

    #[test]
    fn test_weighted_selection() {
        let mut map = CrushMap::new();

        // Create uneven weights
        let root_id = -1;
        let host_id = -2;

        let mut host = CrushBucket::new(host_id, "host0", BucketType::Host);
        host.add_child(0, 0x10000); // weight 1.0
        host.add_child(1, 0x10000 * 3); // weight 3.0

        let mut root = CrushBucket::new(root_id, "root", BucketType::Root);
        root.add_child(host_id, host.weight);

        map.add_bucket(host);
        map.add_bucket(root);
        map.set_root(root_id);
        map.set_weight(0, 1.0);
        map.set_weight(1, 3.0);

        let rule = CrushRule::replicated(0, "replicated_rule", root_id, BucketType::Host);
        map.add_rule(rule);

        // Count selections
        let mut counts = [0u32; 2];
        for pgid in 0..1000u64 {
            let osds = map.select("replicated_rule", pgid, 1).unwrap();
            if !osds.is_empty() {
                counts[osds[0] as usize] += 1;
            }
        }

        // OSD 1 should get roughly 3x more than OSD 0
        // Allow significant variance due to hash distribution
        assert!(
            counts[1] > counts[0],
            "Weighted OSD should be selected more often"
        );
    }

    #[test]
    fn test_multi_host_crush_map() {
        let mut map = CrushMap::new();

        let root_id = -1;
        let host1_id = -2;
        let host2_id = -3;

        // Host 1 with OSDs 0, 1
        let mut host1 = CrushBucket::new(host1_id, "host1", BucketType::Host);
        host1.add_child(0, 0x10000);
        host1.add_child(1, 0x10000);

        // Host 2 with OSDs 2, 3
        let mut host2 = CrushBucket::new(host2_id, "host2", BucketType::Host);
        host2.add_child(2, 0x10000);
        host2.add_child(3, 0x10000);

        // Root containing both hosts
        let mut root = CrushBucket::new(root_id, "root", BucketType::Root);
        root.add_child(host1_id, host1.weight);
        root.add_child(host2_id, host2.weight);

        map.add_bucket(host1);
        map.add_bucket(host2);
        map.add_bucket(root);
        map.set_root(root_id);

        for i in 0..4 {
            map.set_weight(i, 1.0);
        }

        let rule = CrushRule::replicated(0, "replicated_rule", root_id, BucketType::Host);
        map.add_rule(rule);

        let osds = map.select("replicated_rule", 42, 2).unwrap();
        assert_eq!(osds.len(), 2);
    }

    #[test]
    fn test_get_failure_domain() {
        let mut map = CrushMap::new();

        let root_id = -1;
        let host1_id = -2;
        let host2_id = -3;

        let mut host1 = CrushBucket::new(host1_id, "host1", BucketType::Host);
        host1.add_child(0, 0x10000);
        host1.add_child(1, 0x10000);

        let mut host2 = CrushBucket::new(host2_id, "host2", BucketType::Host);
        host2.add_child(2, 0x10000);

        let mut root = CrushBucket::new(root_id, "root", BucketType::Root);
        root.add_child(host1_id, host1.weight);
        root.add_child(host2_id, host2.weight);

        map.add_bucket(host1);
        map.add_bucket(host2);
        map.add_bucket(root);

        assert_eq!(map.get_failure_domain(0, BucketType::Host), Some(host1_id));
        assert_eq!(map.get_failure_domain(1, BucketType::Host), Some(host1_id));
        assert_eq!(map.get_failure_domain(2, BucketType::Host), Some(host2_id));
    }

    #[test]
    fn test_count_osds() {
        let map = CrushMap::simple(8);

        let count = map.count_osds_in_bucket(map.root_id);
        assert_eq!(count, 8);
    }

    #[test]
    fn test_straw2_stability() {
        // Adding an OSD should not massively redistribute existing PGs
        let map1 = CrushMap::simple(7);
        let map2 = CrushMap::simple(8);

        let mut same_count = 0;
        let total_pgs = 100;

        for pgid in 0..total_pgs {
            let osds1 = map1.select("replicated_rule", pgid, 1).unwrap();
            let osds2 = map2.select("replicated_rule", pgid, 1).unwrap();

            if !osds1.is_empty() && !osds2.is_empty() && osds1[0] == osds2[0] {
                same_count += 1;
            }
        }

        // Most PGs should stay on the same OSD (high stability)
        // With 8 OSDs, ~87.5% should stay the same (7/8)
        let stability = (same_count as f64) / (total_pgs as f64);
        assert!(stability > 0.7, "Straw2 stability {} is too low", stability);
    }
}