//! Dynamic shard mapping and scaling.
//!
//! This module implements dynamic shard scaling following Approach B from
//! DYNAMIC_SHARD_SCALING.md: increasing the number of shards and rebalancing
//! producers across them, while maintaining SPSC (single-producer single-consumer)
//! semantics per shard.
//!
//! # Core Principles
//!
//! - Shards never accept multiple producers
//! - Producers get moved to new shards when load increases
//! - Ingestion performance stays at 700M ops/sec per shard
//! - Ordering guarantees remain intact within a shard
//! - Total throughput scales linearly with shard count
//!
//! # Scaling Triggers
//!
//! - Ring buffer fill ratio > threshold (default 70%)
//! - Push latency exceeds threshold (default 5ns)
//! - Batch flush latency exceeds threshold
//! - Session/producer count growth
//!
//! # Architecture
//!
//! ```text
//! Producers -----+
//!                |
//!                v
//!     +----------------------------+
//!     |   Dynamic Shard Mapper     |
//!     +----------------------------+
//!        |         |         |
//!        v         v         v
//!     Shard 0   Shard 1   Shard 2 …
//! ```

use std::sync::atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering as AtomicOrdering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use parking_lot::RwLock;

use crate::config::ScalingPolicy;

/// Metrics for a single shard used for scaling decisions.
#[derive(Debug, Clone)]
pub struct ShardMetrics {
    /// Shard identifier.
    pub shard_id: u16,
    /// Current fill ratio (0.0 - 1.0).
    pub fill_ratio: f64,
    /// Events ingested in the current window.
    pub event_rate: u64,
    /// Average push latency in nanoseconds.
    pub avg_push_latency_ns: u64,
    /// Average batch flush latency in microseconds.
    pub avg_flush_latency_us: u64,
    /// Whether this shard is in drain mode.
    pub draining: bool,
    /// Computed weight for load balancing (lower = less loaded).
    pub weight: f64,
    /// Last update timestamp.
    pub last_updated: Instant,
}

impl ShardMetrics {
    /// Create new metrics for a shard.
    pub fn new(shard_id: u16) -> Self {
        Self {
            shard_id,
            fill_ratio: 0.0,
            event_rate: 0,
            avg_push_latency_ns: 0,
            avg_flush_latency_us: 0,
            draining: false,
            weight: 0.0,
            last_updated: Instant::now(),
        }
    }

    /// Compute the weight based on current metrics.
    /// Lower weight = better candidate for new producers.
    pub fn compute_weight(&mut self) {
        // Weight formula: combines fill ratio, latency, and event rate
        // Higher fill ratio = higher weight (avoid overloaded shards)
        // Higher latency = higher weight
        // Higher event rate = higher weight
        let fill_weight = self.fill_ratio * 100.0;
        let latency_weight = (self.avg_push_latency_ns as f64) / 10.0;
        let rate_weight = (self.event_rate as f64) / 1_000_000.0;

        self.weight = fill_weight + latency_weight + rate_weight;

        // Draining shards get maximum weight (never assign new producers)
        if self.draining {
            self.weight = f64::MAX;
        }
    }
}

/// Live metrics collector for a shard (atomics for lock-free updates).
#[derive(Debug)]
pub struct ShardMetricsCollector {
    /// Shard identifier.
    shard_id: u16,
    /// Ring buffer capacity.
    capacity: usize,
    /// Current buffer length (updated by shard).
    current_len: AtomicU64,
    /// Events ingested in current window.
    events_in_window: AtomicU64,
    /// Sum of push latencies in current window (ns).
    push_latency_sum_ns: AtomicU64,
    /// Number of push operations in current window.
    push_count: AtomicU64,
    /// Sum of flush latencies in current window (us).
    flush_latency_sum_us: AtomicU64,
    /// Number of flush operations in current window.
    flush_count: AtomicU64,
    /// Whether this shard is draining.
    draining: AtomicBool,
    /// Window start time.
    window_start: RwLock<Instant>,
    /// Pushes observed since `set_draining(true)` was last called.
    /// Distinct from `events_in_window` because this counter is NOT
    /// reset by `collect_and_reset`. `finalize_draining` reads this
    /// instead of `events_in_window` so a drain-window-overlap with
    /// a metrics tick can no longer race the counter to zero before
    /// the producer is observed.
    pushes_since_drain_start: AtomicU64,
}

impl ShardMetricsCollector {
    /// Create a new metrics collector.
    pub fn new(shard_id: u16, capacity: usize) -> Self {
        Self {
            shard_id,
            capacity,
            current_len: AtomicU64::new(0),
            events_in_window: AtomicU64::new(0),
            push_latency_sum_ns: AtomicU64::new(0),
            push_count: AtomicU64::new(0),
            flush_latency_sum_us: AtomicU64::new(0),
            flush_count: AtomicU64::new(0),
            draining: AtomicBool::new(false),
            window_start: RwLock::new(Instant::now()),
            pushes_since_drain_start: AtomicU64::new(0),
        }
    }

    /// Record current buffer length.
    #[inline]
    pub fn record_buffer_len(&self, len: usize) {
        self.current_len.store(len as u64, AtomicOrdering::Relaxed);
    }

    /// Record an event ingestion.
    #[inline]
    pub fn record_push(&self, latency_ns: u64) {
        self.events_in_window.fetch_add(1, AtomicOrdering::Relaxed);
        self.push_latency_sum_ns
            .fetch_add(latency_ns, AtomicOrdering::Relaxed);
        self.push_count.fetch_add(1, AtomicOrdering::Relaxed);
        // Always increment — the cost is one fetch_add and the
        // counter only matters when the shard is draining. Cheaper
        // than branching on `self.draining.load()` in the hot path.
        self.pushes_since_drain_start
            .fetch_add(1, AtomicOrdering::Relaxed);
    }

    /// Record a batch flush.
    #[inline]
    pub fn record_flush(&self, latency_us: u64) {
        self.flush_latency_sum_us
            .fetch_add(latency_us, AtomicOrdering::Relaxed);
        self.flush_count.fetch_add(1, AtomicOrdering::Relaxed);
    }

    /// Set drain mode.
    ///
    /// Transitions to draining reset `pushes_since_drain_start` so
    /// `finalize_draining` only counts pushes that arrived after the
    /// drain began.
    ///
    /// A concurrent `record_push` can interleave with the store-zero
    /// on `pushes_since_drain_start` and leave the counter at `1`
    /// rather than `0`. That's not a correctness bug — it just defers
    /// finalization of this shard by one metrics tick — but the
    /// previous code did the store-zero on the counter *before*
    /// publishing the `draining=true` flag, which made the window
    /// slightly larger than necessary. Publishing the flag first
    /// means any push that *observes* `draining=true` is naturally
    /// sequenced after the reset; pushes that beat the flag publish
    /// race the reset just like before.
    ///
    /// We also use `SeqCst` for the publish to give the rest of the
    /// crate a single total order on draining transitions, which
    /// matches the ordering on `try_enter_ingest`'s shutdown flag.
    pub fn set_draining(&self, draining: bool) {
        if draining {
            // Store-zero first (so the flag publish below acts as
            // the release-fence for both writes).
            self.pushes_since_drain_start
                .store(0, AtomicOrdering::SeqCst);
        }
        self.draining.store(draining, AtomicOrdering::SeqCst);
    }

    /// Number of pushes observed since `set_draining(true)` was
    /// last called. Used by `finalize_draining` to detect lingering
    /// producers that the window-reset `events_in_window` counter
    /// can race past.
    pub fn pushes_since_drain_start(&self) -> u64 {
        self.pushes_since_drain_start.load(AtomicOrdering::Relaxed)
    }

    /// Check if draining.
    pub fn is_draining(&self) -> bool {
        self.draining.load(AtomicOrdering::Acquire)
    }

    /// Collect metrics and reset window counters.
    ///
    /// NOTE: The individual atomic swaps below are not collectively atomic with
    /// respect to concurrent `record_push`/`record_flush` calls. This means a
    /// push recorded between, say, the `events_in_window` swap and the
    /// `push_count` swap could be counted in one counter but not the other for
    /// a given window. This small inaccuracy is an accepted trade-off to
    /// preserve the lock-free design of the hot path (`record_push` /
    /// `record_flush`). Adding a `Mutex` here would serialize the hot path
    /// and defeat the purpose of using atomics.
    pub fn collect_and_reset(&self) -> ShardMetrics {
        let current_len = self.current_len.load(AtomicOrdering::Relaxed);
        let events = self.events_in_window.swap(0, AtomicOrdering::Relaxed);
        let push_latency_sum = self.push_latency_sum_ns.swap(0, AtomicOrdering::Relaxed);
        let push_count = self.push_count.swap(0, AtomicOrdering::Relaxed);
        let flush_latency_sum = self.flush_latency_sum_us.swap(0, AtomicOrdering::Relaxed);
        let flush_count = self.flush_count.swap(0, AtomicOrdering::Relaxed);

        let fill_ratio = if self.capacity > 0 {
            current_len as f64 / self.capacity as f64
        } else {
            0.0
        };

        let avg_push_latency = push_latency_sum.checked_div(push_count).unwrap_or(0);

        let avg_flush_latency = flush_latency_sum.checked_div(flush_count).unwrap_or(0);

        // Reset window
        *self.window_start.write() = Instant::now();

        let mut metrics = ShardMetrics {
            shard_id: self.shard_id,
            fill_ratio,
            event_rate: events,
            avg_push_latency_ns: avg_push_latency,
            avg_flush_latency_us: avg_flush_latency,
            draining: self.draining.load(AtomicOrdering::Acquire),
            weight: 0.0,
            last_updated: Instant::now(),
        };
        metrics.compute_weight();
        metrics
    }
}

/// Scaling decision made by the mapper.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ScalingDecision {
    /// No scaling needed.
    None,
    /// Scale up by adding N shards.
    ScaleUp(u16),
    /// Scale down by removing N shards (marks them for draining).
    ScaleDown(u16),
}

/// Errors that can occur during scaling operations.
#[derive(Debug, Clone)]
pub enum ScalingError {
    /// Invalid scaling policy.
    InvalidPolicy(String),
    /// Already at maximum shards.
    AtMaxShards,
    /// Already at minimum shards.
    AtMinShards,
    /// Scaling operation in cooldown.
    InCooldown,
    /// Shard creation failed.
    ShardCreationFailed(String),
}

impl std::fmt::Display for ScalingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidPolicy(msg) => write!(f, "invalid scaling policy: {}", msg),
            Self::AtMaxShards => write!(f, "already at maximum shard count"),
            Self::AtMinShards => write!(f, "already at minimum shard count"),
            Self::InCooldown => write!(f, "scaling operation in cooldown"),
            Self::ShardCreationFailed(msg) => write!(f, "shard creation failed: {}", msg),
        }
    }
}

impl std::error::Error for ScalingError {}

/// State of a shard in the mapper.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShardState {
    /// Shard is being provisioned: id allocated and metrics collector
    /// in place, but `select_shard` must not route to it yet because
    /// upstream workers (drain / batch) have not been spawned. Caller
    /// transitions to `Active` via `activate` once the workers are
    /// ready. Closes the race where a freshly-added shard accepted
    /// producer pushes before any consumer existed.
    Provisioning,
    /// Shard is active and accepting producers.
    Active,
    /// Shard is draining (no new producers, waiting for empty).
    Draining,
    /// Shard is stopped and can be removed.
    Stopped,
}

/// Information about a shard managed by the mapper.
#[derive(Debug)]
struct MappedShard {
    /// Shard ID.
    id: u16,
    /// Current state.
    state: ShardState,
    /// Metrics collector.
    metrics: Arc<ShardMetricsCollector>,
    /// When this shard entered drain mode (if draining).
    drain_started: Option<Instant>,
    /// Last collected metrics.
    last_metrics: ShardMetrics,
}

/// Callback type for shard lifecycle events.
type ShardCallback = Box<dyn Fn(u16) + Send + Sync>;

/// Dynamic shard mapper that manages shard allocation and producer routing.
///
/// This is the core component for dynamic scaling. It:
/// - Tracks metrics for all shards
/// - Makes scaling decisions based on policy
/// - Routes producers to the least-loaded shards
/// - Manages shard lifecycle (active → draining → stopped)
pub struct ShardMapper {
    /// Mapped shards (RwLock for concurrent reads, rare writes).
    shards: RwLock<Vec<MappedShard>>,
    /// Current active shard count.
    active_count: AtomicU16,
    /// Scaling policy.
    ///
    /// This field is immutable for the lifetime of the mapper. The
    /// previous `set_policy(&mut self, …)` API was unreachable in
    /// practice — every production callsite holds the mapper behind
    /// an `Arc`, and `Arc::get_mut` requires a strong count of 1,
    /// which never holds once the worker pool clones the `Arc`. The
    /// method has been removed; recreate the mapper (and
    /// the bus that owns it) to change the policy.
    policy: ScalingPolicy,
    /// Ring buffer capacity for new shards.
    ring_buffer_capacity: usize,
    /// Last scaling operation timestamp.
    last_scaling: RwLock<Option<Instant>>,
    /// Callback for shard creation (provided by ShardManager).
    on_shard_created: RwLock<Option<ShardCallback>>,
    /// Callback for shard removal (provided by ShardManager).
    on_shard_removed: RwLock<Option<ShardCallback>>,
    /// Monotonic shard-id allocator. The next `scale_up` call gets
    /// `fetch_add(1)` from here. Distinct from `shards.iter().max() +
    /// 1`: that approach reused ids whenever the highest-numbered
    /// shard had been drained-and-removed, silently merging two
    /// unrelated shard lifetimes in any external system that keys
    /// metrics or checkpoints on shard id. Monotonic allocation
    /// ensures every shard ever allocated has a globally unique id
    /// for the lifetime of this mapper.
    next_shard_id: AtomicU16,
}

impl ShardMapper {
    /// Create a new shard mapper with the given initial shard count and policy.
    pub fn new(
        initial_shards: u16,
        ring_buffer_capacity: usize,
        policy: ScalingPolicy,
    ) -> Result<Self, ScalingError> {
        let mut policy = policy.normalize();
        // Ensure max_shards can accommodate the initial shard count
        if policy.max_shards < initial_shards {
            policy.max_shards = initial_shards;
        }
        policy
            .validate()
            .map_err(|e| ScalingError::InvalidPolicy(e.to_string()))?;

        let shards: Vec<MappedShard> = (0..initial_shards)
            .map(|id| MappedShard {
                id,
                state: ShardState::Active,
                metrics: Arc::new(ShardMetricsCollector::new(id, ring_buffer_capacity)),
                drain_started: None,
                last_metrics: ShardMetrics::new(id),
            })
            .collect();

        Ok(Self {
            shards: RwLock::new(shards),
            active_count: AtomicU16::new(initial_shards),
            policy,
            ring_buffer_capacity,
            last_scaling: RwLock::new(None),
            on_shard_created: RwLock::new(None),
            on_shard_removed: RwLock::new(None),
            // Initial shards occupy ids `[0, initial_shards)`, so the
            // first scale-up takes `initial_shards`.
            next_shard_id: AtomicU16::new(initial_shards),
        })
    }

    /// Set callback for shard creation.
    pub fn set_on_shard_created<F>(&self, callback: F)
    where
        F: Fn(u16) + Send + Sync + 'static,
    {
        *self.on_shard_created.write() = Some(Box::new(callback));
    }

    /// Set callback for shard removal.
    pub fn set_on_shard_removed<F>(&self, callback: F)
    where
        F: Fn(u16) + Send + Sync + 'static,
    {
        *self.on_shard_removed.write() = Some(Box::new(callback));
    }

    /// Get the metrics collector for a shard.
    pub fn metrics_collector(&self, shard_id: u16) -> Option<Arc<ShardMetricsCollector>> {
        let shards = self.shards.read();
        shards
            .iter()
            .find(|s| s.id == shard_id)
            .map(|s| s.metrics.clone())
    }

    /// Get current active shard count.
    pub fn active_shard_count(&self) -> u16 {
        self.active_count.load(AtomicOrdering::Acquire)
    }

    /// Get total shard count (including draining).
    pub fn total_shard_count(&self) -> u16 {
        self.shards.read().len() as u16
    }

    /// Select the best shard for a new event/producer.
    ///
    /// This implements weighted shard selection:
    /// - Only considers active (non-draining) shards
    /// - Prefers shards with lower weight (less loaded)
    /// - Falls back to round-robin if weights are equal
    #[inline]
    pub fn select_shard(&self, event_hash: u64) -> u16 {
        let shards = self.shards.read();

        // Filter to active shards only
        let active: Vec<_> = shards
            .iter()
            .filter(|s| s.state == ShardState::Active)
            .collect();

        if active.is_empty() {
            // Previously fell back to `find(|s| s.state !=
            // ShardState::Stopped)`, which silently routed to a
            // `Draining` shard. Pushes to a draining shard increment
            // `pushes_since_drain_start`, blocking finalization
            // indefinitely.
            //
            // `Provisioning` shards aren't routable either — they
            // exist in the mapper but the routing table doesn't have
            // them yet, so any caller that tries to push lands
            // in `resolve_idx → None` and surfaces as "unrouted" via
            // the manager-level counter. That's the correct signal:
            // the system has no destination for this event, the
            // caller should back off and retry.
            //
            // Return `u16::MAX` (out-of-band sentinel) so callers
            // that look up the id in the routing table get a
            // definite miss, which the manager already accounts for.
            return u16::MAX;
        }

        // Find the shard with lowest weight
        let min_weight = active
            .iter()
            .map(|s| s.last_metrics.weight)
            .fold(f64::MAX, f64::min);

        // Get all shards with the minimum weight (within tolerance)
        let candidates: Vec<_> = active
            .iter()
            .filter(|s| (s.last_metrics.weight - min_weight).abs() < 0.1)
            .collect();

        // Use hash to pick among candidates for determinism
        // Fallback to first active shard if tolerance filter excludes all (e.g. NaN weights)
        if candidates.is_empty() {
            return active[0].id;
        }
        let idx = (event_hash as usize) % candidates.len();
        candidates[idx].id
    }

    /// Collect metrics from all shards and update weights.
    pub fn collect_metrics(&self) -> Vec<ShardMetrics> {
        let mut shards = self.shards.write();
        shards
            .iter_mut()
            .map(|s| {
                s.last_metrics = s.metrics.collect_and_reset();
                s.last_metrics.clone()
            })
            .collect()
    }

    /// Evaluate scaling based on current metrics.
    ///
    /// Returns a scaling decision without executing it.
    pub fn evaluate_scaling(&self) -> ScalingDecision {
        if !self.policy.auto_scale {
            return ScalingDecision::None;
        }

        // Check cooldown
        if let Some(last) = *self.last_scaling.read() {
            if last.elapsed() < self.policy.cooldown {
                return ScalingDecision::None;
            }
        }

        let shards = self.shards.read();
        let active_count = self.active_count.load(AtomicOrdering::Acquire);

        // Check for scale-up triggers
        let mut overloaded_count = 0;
        let mut underutilized_count = 0;

        for shard in shards.iter() {
            if shard.state != ShardState::Active {
                continue;
            }

            let m = &shard.last_metrics;

            // Scale-up triggers
            if m.fill_ratio > self.policy.fill_ratio_threshold
                || m.avg_push_latency_ns > self.policy.push_latency_threshold_ns
                || m.avg_flush_latency_us > self.policy.flush_latency_threshold_us
            {
                overloaded_count += 1;
            }

            // Scale-down triggers
            if m.fill_ratio < self.policy.underutilized_threshold && m.event_rate == 0 {
                underutilized_count += 1;
            }
        }

        // Scale up if more than half of shards are overloaded
        if overloaded_count > active_count / 2 && active_count < self.policy.max_shards {
            // Add shards proportional to overload
            let shards_to_add = (overloaded_count / 2)
                .max(1)
                .min(self.policy.max_shards - active_count);
            return ScalingDecision::ScaleUp(shards_to_add);
        }

        // Scale down if more than half of shards are underutilized
        // and we're above minimum
        if underutilized_count > active_count / 2 && active_count > self.policy.min_shards {
            let shards_to_remove = (underutilized_count / 2)
                .max(1)
                .min(active_count - self.policy.min_shards);
            return ScalingDecision::ScaleDown(shards_to_remove);
        }

        ScalingDecision::None
    }

    /// Validate cooldown + max-shards budget for an `add count`
    /// scale-up request. Cheap pre-check that doesn't take the
    /// write lock — callers re-check under the lock.
    fn check_scale_up_budget(&self, count: u16) -> Result<(), ScalingError> {
        let current = self.active_count.load(AtomicOrdering::Acquire);
        let would_be = current
            .checked_add(count)
            .ok_or(ScalingError::AtMaxShards)?;
        if would_be > self.policy.max_shards {
            return Err(ScalingError::AtMaxShards);
        }
        let last = self.last_scaling.read();
        if let Some(ts) = *last {
            if ts.elapsed() < self.policy.cooldown {
                return Err(ScalingError::InCooldown);
            }
        }
        Ok(())
    }

    /// Allocate `count` shard ids and push their `MappedShard`
    /// records into `shards` with the supplied `state`. The caller
    /// already holds `self.shards.write()` and is responsible for
    /// dropping it before notifying callbacks. Returns the allocated
    /// ids in order.
    ///
    /// Performs the budget + cooldown re-check under the write lock,
    /// the next_shard_id allocation, and the per-shard push. Does NOT
    /// touch `active_count` — `scale_up` bumps it for `Active` shards;
    /// `Provisioning` shards bump it later when `activate` fires.
    fn allocate_shards_inner(
        &self,
        count: u16,
        state: ShardState,
        shards: &mut Vec<MappedShard>,
    ) -> Result<Vec<u16>, ScalingError> {
        // Re-check budget under the write lock — two concurrent
        // scale-up callers could both pass the read-locked early
        // check, both serialize through `shards.write()`, and both
        // succeed without this re-check.
        self.check_scale_up_budget(count)?;

        let first_id = self.next_shard_id.load(AtomicOrdering::Relaxed);
        let last_needed = first_id
            .checked_add(count.saturating_sub(1))
            .ok_or(ScalingError::AtMaxShards)?;
        // Reserve `u16::MAX` as a sentinel so the post-allocation
        // store cannot wrap.
        if last_needed == u16::MAX {
            return Err(ScalingError::AtMaxShards);
        }
        self.next_shard_id
            .store(first_id + count, AtomicOrdering::Relaxed);

        let mut new_ids = Vec::with_capacity(count as usize);
        for i in 0..count {
            let new_id = first_id + i;
            shards.push(MappedShard {
                id: new_id,
                state,
                metrics: Arc::new(ShardMetricsCollector::new(
                    new_id,
                    self.ring_buffer_capacity,
                )),
                drain_started: None,
                last_metrics: ShardMetrics::new(new_id),
            });
            new_ids.push(new_id);
        }
        Ok(new_ids)
    }

    /// Execute a scale-up operation.
    ///
    /// Creates new shards in the `Active` state and makes them
    /// immediately available for routing. Use [`scale_up_provisioning`]
    /// if upstream workers (drain / batch) need to be wired up before
    /// the shard becomes selectable — otherwise producer pushes can
    /// race ahead of consumer creation.
    ///
    /// [`scale_up_provisioning`]: Self::scale_up_provisioning
    pub fn scale_up(&self, count: u16) -> Result<Vec<u16>, ScalingError> {
        // Short-circuit `count == 0` so a no-op call doesn't bump the
        // cooldown timestamp or trip the `u16::MAX` sentinel check
        // (which previously fired spuriously when
        // `first_id == u16::MAX` even though zero ids were being
        // allocated).
        if count == 0 {
            return Ok(Vec::new());
        }

        self.check_scale_up_budget(count)?;

        let mut shards = self.shards.write();
        let new_ids = self.allocate_shards_inner(count, ShardState::Active, &mut shards)?;

        // Update counts. `fetch_add` cannot wrap here — the
        // `check_scale_up_budget` gate above ensures `current + count <=
        // max_shards <= u16::MAX` — but keep the ordering explicit
        // so the next contender's cooldown re-check sees the fresh
        // timestamp.
        self.active_count.fetch_add(count, AtomicOrdering::Release);
        *self.last_scaling.write() = Some(Instant::now());

        // Drop the write lock before notifying callbacks — they
        // are user-supplied and may take arbitrary time.
        drop(shards);

        if let Some(callback) = self.on_shard_created.read().as_ref() {
            for &id in &new_ids {
                callback(id);
            }
        }

        Ok(new_ids)
    }

    /// Like [`scale_up`], but the new shards are created in the
    /// `Provisioning` state. They receive an id and a metrics
    /// collector, but `select_shard` will not route to them and they
    /// are excluded from `active_shard_count` / `evaluate_scaling`
    /// until the caller transitions each shard with [`activate`].
    ///
    /// Use this when upstream consumer infrastructure (drain/batch
    /// workers, mpsc channels, etc.) must be wired up *before* the
    /// shard becomes selectable. Without this gating, producers can
    /// observe the shard via `select_shard`, push into its ring
    /// buffer, and never have those events drained.
    ///
    /// Returns the allocated ids in order. Cooldown / `max_shards`
    /// gating matches `scale_up` so that staged allocation cannot be
    /// used to bypass the policy.
    ///
    /// [`scale_up`]: Self::scale_up
    /// [`activate`]: Self::activate
    pub fn scale_up_provisioning(&self, count: u16) -> Result<Vec<u16>, ScalingError> {
        if count == 0 {
            return Ok(Vec::new());
        }

        self.check_scale_up_budget(count)?;

        let mut shards = self.shards.write();
        let new_ids = self.allocate_shards_inner(count, ShardState::Provisioning, &mut shards)?;

        // Bump cooldown but NOT active_count — these shards are
        // not active yet. `activate` bumps active_count when each
        // becomes selectable.
        *self.last_scaling.write() = Some(Instant::now());
        drop(shards);
        // Intentionally do NOT fire `on_shard_created` here — the
        // callback signals "shard is live"; provisioning shards are
        // not. `activate` fires it instead.
        Ok(new_ids)
    }

    /// Transition a `Provisioning` shard to `Active`.
    ///
    /// Idempotent for `Active` shards (returns `Ok(())` without
    /// re-firing the callback). Returns `InvalidPolicy` for unknown
    /// or `Draining`/`Stopped` shards — those states require a
    /// different lifecycle path. Bumps `active_count` and notifies
    /// the `on_shard_created` callback exactly once.
    pub fn activate(&self, shard_id: u16) -> Result<(), ScalingError> {
        let mut shards = self.shards.write();
        let shard = shards
            .iter_mut()
            .find(|s| s.id == shard_id)
            .ok_or_else(|| {
                ScalingError::InvalidPolicy(format!("activate: shard {} not found", shard_id))
            })?;
        match shard.state {
            ShardState::Active => return Ok(()),
            ShardState::Provisioning => {
                // Gate activation on
                // `active_count < max_shards`. The budget gate at
                // `check_scale_up_budget` only counts ALREADY-active
                // shards — multiple `scale_up_provisioning(1)` calls
                // can each pass (they don't bump `active_count`),
                // and an unconditional `fetch_add(1)` here would
                // push past `max_shards`. Subsequent
                // `evaluate_scaling`'s `max_shards - active_count`
                // arithmetic would then underflow u16 (debug-build
                // panic; release wraps to ~65530). Activate-time
                // re-checks the budget with the lock held; the
                // Provisioning shard stays in Provisioning state and
                // the caller (e.g. `add_shard_internal`'s rollback)
                // is responsible for tearing it down.
                let current = self.active_count.load(AtomicOrdering::Acquire);
                if current >= self.policy.max_shards {
                    return Err(ScalingError::AtMaxShards);
                }
                shard.state = ShardState::Active;
            }
            ShardState::Draining | ShardState::Stopped => {
                return Err(ScalingError::InvalidPolicy(format!(
                    "activate: shard {} is in state {:?}, cannot activate",
                    shard_id, shard.state
                )));
            }
        }
        drop(shards);

        // Bump active_count to mirror what `scale_up` would have done.
        // The budget gate above ensures we can never push past
        // `max_shards`.
        self.active_count.fetch_add(1, AtomicOrdering::Release);

        if let Some(callback) = self.on_shard_created.read().as_ref() {
            callback(shard_id);
        }
        Ok(())
    }

    /// Drain a specific shard by id, transitioning it from `Active`
    /// to `Draining`.
    ///
    /// Companion to `ShardManager::drain_shard`. The previous
    /// implementation only flipped the metrics collector's `draining`
    /// atomic; this version atomically updates `MappedShard.state`
    /// (so `select_shard` stops routing to the shard) and decrements
    /// `active_count` (so `evaluate_scaling`'s budget math stays
    /// consistent with `scale_down`). Returns an error if the shard
    /// is not in `Active` state, or if doing so would push the active
    /// count below `min_shards`.
    pub fn drain_specific(&self, shard_id: u16) -> Result<(), ScalingError> {
        let mut shards = self.shards.write();
        let current_active = self.active_count.load(AtomicOrdering::Acquire);
        if current_active <= self.policy.min_shards {
            return Err(ScalingError::AtMinShards);
        }
        let shard = shards
            .iter_mut()
            .find(|s| s.id == shard_id)
            .ok_or_else(|| {
                ScalingError::InvalidPolicy(format!("drain_specific: shard {} not found", shard_id))
            })?;
        match shard.state {
            ShardState::Active => {
                shard.state = ShardState::Draining;
                shard.drain_started = Some(Instant::now());
                shard.metrics.set_draining(true);
            }
            ShardState::Draining => return Ok(()),
            ShardState::Provisioning | ShardState::Stopped => {
                return Err(ScalingError::InvalidPolicy(format!(
                    "drain_specific: shard {} is in state {:?}, cannot drain",
                    shard_id, shard.state
                )));
            }
        }
        drop(shards);
        self.active_count.fetch_sub(1, AtomicOrdering::Release);
        Ok(())
    }

    /// Start draining shards for scale-down.
    ///
    /// Marks shards as draining so they stop receiving new events.
    /// Shards will be removed once they're empty.
    pub fn scale_down(&self, count: u16) -> Result<Vec<u16>, ScalingError> {
        // Early checks (may race, but avoid acquiring the write lock)
        let current = self.active_count.load(AtomicOrdering::Acquire);
        if current <= self.policy.min_shards {
            return Err(ScalingError::AtMinShards);
        }

        let to_drain = count.min(current - self.policy.min_shards);
        if to_drain == 0 {
            return Err(ScalingError::AtMinShards);
        }
        {
            let last = self.last_scaling.read();
            if let Some(ts) = *last {
                if ts.elapsed() < self.policy.cooldown {
                    return Err(ScalingError::InCooldown);
                }
            }
        }

        let mut shards = self.shards.write();

        // Re-check under the lock to prevent race conditions (double-check pattern)
        let current = self.active_count.load(AtomicOrdering::Acquire);
        if current <= self.policy.min_shards {
            return Err(ScalingError::AtMinShards);
        }
        let to_drain = count.min(current - self.policy.min_shards);
        if to_drain == 0 {
            return Err(ScalingError::AtMinShards);
        }
        // Re-check cooldown under the same write lock that gates
        // mutation — see the matching note in `scale_up`.
        {
            let last = self.last_scaling.read();
            if let Some(ts) = *last {
                if ts.elapsed() < self.policy.cooldown {
                    return Err(ScalingError::InCooldown);
                }
            }
        }

        let mut drained_ids = Vec::with_capacity(to_drain as usize);

        // Find shards with lowest weight (least utilized) to drain
        let mut active_indices: Vec<_> = shards
            .iter()
            .enumerate()
            .filter(|(_, s)| s.state == ShardState::Active)
            .map(|(i, s)| (i, s.last_metrics.weight))
            .collect();

        // Sort by weight (ascending - drain least utilized first)
        active_indices.sort_by(|a, b| a.1.total_cmp(&b.1));

        // Mark shards for draining
        for (idx, _) in active_indices.into_iter().take(to_drain as usize) {
            shards[idx].state = ShardState::Draining;
            shards[idx].drain_started = Some(Instant::now());
            shards[idx].metrics.set_draining(true);
            drained_ids.push(shards[idx].id);
        }

        // Update count
        self.active_count
            .fetch_sub(to_drain, AtomicOrdering::Release);
        *self.last_scaling.write() = Some(Instant::now());

        Ok(drained_ids)
    }

    /// Check draining shards and finalize those that are empty.
    ///
    /// Returns IDs of shards that were stopped.
    ///
    /// This predicate looks ONLY at the ring buffer
    /// (`current_len` + `pushes_since_drain_start`); it does NOT
    /// probe the per-shard mpsc channel or the BatchWorker's
    /// `current_batch`. A shard that the predicate flags as empty
    /// can still have events queued in those two places. The
    /// correctness gate is therefore `bus::remove_shard_internal`,
    /// which awaits the BatchWorker's `JoinHandle` before
    /// constructing the stranded-flush batch — see that function's
    /// step 3 for the rationale. Tightening this predicate is a
    /// defense-in-depth follow-up; a stricter ring-buffer-empty
    /// signal here would only narrow an already-closed window.
    pub fn finalize_draining(&self) -> Vec<u16> {
        let mut shards = self.shards.write();
        let mut stopped = Vec::new();

        for shard in shards.iter_mut() {
            if shard.state == ShardState::Draining {
                // Check if shard is empty by reading current_len directly,
                // avoiding collect_and_reset() which destructively zeros all counters.
                let current_len = shard.metrics.current_len.load(AtomicOrdering::Relaxed);
                let fill_ratio = if shard.metrics.capacity > 0 {
                    current_len as f64 / shard.metrics.capacity as f64
                } else {
                    0.0
                };
                // Previously read `events_in_window` here, which
                // `collect_and_reset` zeros every metrics tick. A
                // producer push that landed in the window between two
                // ticks could be silently zeroed out, so a draining
                // shard whose buffer transiently emptied was finalized
                // with a producer still attached.
                // `pushes_since_drain_start` is a separate counter
                // that is only reset by `set_draining(true)`, so any
                // push observed since the drain began is sticky —
                // exactly the signal we want.
                let pushes_after_drain = shard
                    .metrics
                    .pushes_since_drain_start
                    .load(AtomicOrdering::Relaxed);
                if fill_ratio == 0.0 && pushes_after_drain == 0 {
                    // Check if we've waited long enough
                    if let Some(drain_start) = shard.drain_started {
                        if drain_start.elapsed() > Duration::from_millis(100) {
                            shard.state = ShardState::Stopped;
                            stopped.push(shard.id);
                        }
                    }
                }
            }
        }

        // Drop the write lock BEFORE notifying. The callback is
        // user-supplied and may re-enter the mapper (`shard_state`,
        // `select_shard`, `metrics_collector`, …), each of which
        // acquires `shards.read()`. `parking_lot::RwLock` is not
        // recursive, so a re-entrant read attempt while we hold a
        // write would deadlock. `scale_up`'s callback path already
        // releases its lock before calling out — mirror that
        // here.
        drop(shards);

        if !stopped.is_empty() {
            if let Some(callback) = self.on_shard_removed.read().as_ref() {
                for &id in &stopped {
                    callback(id);
                }
            }
        }

        stopped
    }

    /// Remove stopped shards from the mapper.
    pub fn remove_stopped_shards(&self) -> Vec<u16> {
        let mut shards = self.shards.write();
        let before = shards.len();
        let removed: Vec<u16> = shards
            .iter()
            .filter(|s| s.state == ShardState::Stopped)
            .map(|s| s.id)
            .collect();

        shards.retain(|s| s.state != ShardState::Stopped);

        if shards.len() < before {
            tracing::info!(
                removed = removed.len(),
                remaining = shards.len(),
                "Removed stopped shards"
            );
        }

        removed
    }

    /// Get the state of a specific shard.
    pub fn shard_state(&self, shard_id: u16) -> Option<ShardState> {
        self.shards
            .read()
            .iter()
            .find(|s| s.id == shard_id)
            .map(|s| s.state)
    }

    /// Get all active shard IDs.
    pub fn active_shard_ids(&self) -> Vec<u16> {
        self.shards
            .read()
            .iter()
            .filter(|s| s.state == ShardState::Active)
            .map(|s| s.id)
            .collect()
    }

    /// Get all shard IDs (including draining).
    pub fn all_shard_ids(&self) -> Vec<u16> {
        self.shards
            .read()
            .iter()
            .filter(|s| s.state != ShardState::Stopped)
            .map(|s| s.id)
            .collect()
    }

    /// Get the scaling policy.
    pub fn policy(&self) -> &ScalingPolicy {
        &self.policy
    }
    // `set_policy` previously took `&mut self` and was unreachable
    // through the `Arc<ShardMapper>` that the production code holds
    // (`Arc::get_mut` fails once the worker pool has cloned the Arc).
    // The method has been removed — recreate the mapper / bus to
    // change the policy.
}

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

    #[test]
    fn test_shard_mapper_creation() {
        let mapper = ShardMapper::new(4, 1024, ScalingPolicy::default()).unwrap();
        assert_eq!(mapper.active_shard_count(), 4);
        assert_eq!(mapper.total_shard_count(), 4);
    }

    #[test]
    fn test_select_shard_distributes() {
        let mapper = ShardMapper::new(4, 1024, ScalingPolicy::default()).unwrap();

        // Different hashes should potentially select different shards
        let mut selected = std::collections::HashSet::new();
        for i in 0..100u64 {
            let shard = mapper.select_shard(i * 12345);
            selected.insert(shard);
        }

        // With 4 shards, we should hit multiple
        assert!(!selected.is_empty());
    }

    #[test]
    fn test_scale_up() {
        // Explicitly set max_shards to allow scaling from 2 to 4
        let policy = ScalingPolicy {
            max_shards: 8,
            ..Default::default()
        };
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();

        let new_ids = mapper.scale_up(2).unwrap();
        assert_eq!(new_ids.len(), 2);
        assert_eq!(mapper.active_shard_count(), 4);
    }

    #[test]
    fn test_scale_up_max_limit() {
        let policy = ScalingPolicy {
            max_shards: 4,
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        let result = mapper.scale_up(1);
        assert!(matches!(result, Err(ScalingError::AtMaxShards)));
    }

    #[test]
    fn test_scale_down() {
        let policy = ScalingPolicy {
            min_shards: 1,
            cooldown: Duration::from_nanos(1), // Disable cooldown for test
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        let drained = mapper.scale_down(2).unwrap();
        assert_eq!(drained.len(), 2);
        assert_eq!(mapper.active_shard_count(), 2);
    }

    #[test]
    fn test_scale_down_min_limit() {
        let policy = ScalingPolicy {
            min_shards: 4,
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        let result = mapper.scale_down(1);
        assert!(matches!(result, Err(ScalingError::AtMinShards)));
    }

    #[test]
    fn test_metrics_collection() {
        let mapper = ShardMapper::new(2, 1024, ScalingPolicy::default()).unwrap();

        // Record some metrics
        if let Some(collector) = mapper.metrics_collector(0) {
            collector.record_buffer_len(512);
            collector.record_push(5);
            collector.record_push(10);
        }

        let metrics = mapper.collect_metrics();
        assert_eq!(metrics.len(), 2);

        let shard0_metrics = metrics.iter().find(|m| m.shard_id == 0).unwrap();
        assert!(shard0_metrics.fill_ratio > 0.0);
    }

    #[test]
    fn test_draining_excludes_from_selection() {
        let policy = ScalingPolicy {
            min_shards: 1,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();

        // Drain one shard
        let drained = mapper.scale_down(1).unwrap();
        assert_eq!(drained.len(), 1);

        // All selections should go to the remaining active shard
        let active_ids = mapper.active_shard_ids();
        assert_eq!(active_ids.len(), 1);

        for i in 0..100u64 {
            let selected = mapper.select_shard(i);
            assert!(active_ids.contains(&selected));
        }
    }

    #[test]
    fn test_policy_validation() {
        let invalid_policy = ScalingPolicy {
            fill_ratio_threshold: 1.5, // Invalid
            ..Default::default()
        };
        assert!(invalid_policy.validate().is_err());

        // Without normalize(), this would be invalid
        let invalid_policy2 = ScalingPolicy {
            min_shards: 10,
            max_shards: 5,
            ..Default::default()
        };
        assert!(invalid_policy2.validate().is_err());
    }

    #[test]
    fn test_policy_normalize_auto_adjusts_max_shards() {
        // When min_shards > max_shards, normalize() should adjust max_shards
        let policy = ScalingPolicy {
            min_shards: 8,
            max_shards: 2, // Less than min_shards
            ..Default::default()
        };

        let normalized = policy.normalize();
        assert_eq!(
            normalized.max_shards, 8,
            "max_shards should be adjusted to min_shards"
        );
        assert!(
            normalized.validate().is_ok(),
            "normalized policy should be valid"
        );
    }

    /// Regression: BUG_REPORT.md #7 — `scale_up` previously allocated
    /// new shard ids as `shards.iter().max() + 1`, which reused ids
    /// after the highest-numbered shard was drained-and-removed.
    /// Reusing ids merges two distinct shard lifetimes in any
    /// external metric/checkpoint system that keys on shard id.
    /// The fix uses a monotonic `next_shard_id` counter.
    #[test]
    fn scale_up_does_not_reuse_ids_after_remove() {
        let policy = ScalingPolicy {
            min_shards: 1,
            max_shards: 16,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();

        // Initial ids are 0 and 1. Scale up to 4 — new ids must be
        // 2 and 3 (the next two slots from the monotonic counter).
        let new_ids = mapper.scale_up(2).unwrap();
        assert_eq!(new_ids, vec![2, 3]);

        // Drain one shard. `scale_down` picks by lowest weight, so
        // we don't get to choose which id is drained. Whichever it
        // is, we then force it through Stopped + remove and check
        // that the removed id is *not* reissued by the next
        // scale_up.
        let drained = mapper.scale_down(1).unwrap();
        let drained_id = drained[0];
        for shard in mapper.shards.write().iter_mut() {
            if shard.id == drained_id {
                shard.state = ShardState::Stopped;
            }
        }
        let removed = mapper.remove_stopped_shards();
        assert_eq!(removed, vec![drained_id]);

        // Scale up by 1. The broken `max(existing_ids) + 1` allocator
        // could revive the just-removed id whenever `drained_id` had
        // been the highest id present (e.g. id 3 with 0/1/2/3 → 0/1/2
        // → next would be 3 again). The fix uses the monotonic
        // counter, which sits at 4 after the earlier scale_up, so
        // the new id must be 4 regardless of which id was drained.
        let new_ids = mapper.scale_up(1).unwrap();
        assert_eq!(
            new_ids,
            vec![4],
            "shard id {drained_id} was just removed; reusing any \
             previously-issued id would merge two distinct shard \
             lifetimes in external systems"
        );
    }

    #[test]
    fn test_policy_normalize_preserves_valid_config() {
        // When max_shards >= min_shards, normalize() should not change anything
        let policy = ScalingPolicy {
            min_shards: 4,
            max_shards: 16,
            ..Default::default()
        };

        let normalized = policy.normalize();
        assert_eq!(normalized.min_shards, 4);
        assert_eq!(normalized.max_shards, 16);
    }

    #[test]
    fn test_shard_mapper_normalizes_policy() {
        // ShardMapper should accept a policy where min_shards > default max_shards
        // because it calls normalize() internally
        let policy = ScalingPolicy {
            min_shards: 4,
            ..Default::default()
        };

        // This should succeed even on machines with < 4 CPUs
        let result = ShardMapper::new(4, 1024, policy);
        assert!(
            result.is_ok(),
            "ShardMapper should normalize policy automatically"
        );
    }

    #[test]
    fn test_shard_mapper_adjusts_max_shards_to_initial_count() {
        // ShardMapper should adjust max_shards to accommodate initial_shards
        // even if initial_shards > default max_shards (CPU count)
        let policy = ScalingPolicy::default();

        // Create mapper with 8 initial shards - should work even on 2-core machines
        let result = ShardMapper::new(8, 1024, policy);
        assert!(
            result.is_ok(),
            "ShardMapper should adjust max_shards to initial_shards"
        );

        let mapper = result.unwrap();
        assert_eq!(mapper.active_shard_count(), 8);

        // Verify the policy was adjusted
        assert!(
            mapper.policy().max_shards >= 8,
            "max_shards should be at least initial_shards"
        );
    }

    #[test]
    fn test_scale_up_max_shards_concurrent() {
        use std::sync::Arc;
        use std::thread;

        let policy = ScalingPolicy {
            max_shards: 10,
            cooldown: Duration::from_nanos(1), // Disable cooldown for test
            ..Default::default()
        };
        let mapper = Arc::new(ShardMapper::new(5, 1024, policy).unwrap());

        // Spawn multiple threads that all try to scale up
        let mut handles = vec![];
        for _ in 0..5 {
            let mapper_clone = mapper.clone();
            handles.push(thread::spawn(move || mapper_clone.scale_up(3)));
        }

        // Collect results
        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();

        // Some should succeed, some should fail with AtMaxShards
        let successes: Vec<_> = results.iter().filter(|r| r.is_ok()).collect();
        let failures: Vec<_> = results
            .iter()
            .filter(|r| matches!(r, Err(ScalingError::AtMaxShards)))
            .collect();

        // We started with 5, max is 10, each tries to add 3
        // At most we can add 5 more shards, so at most 1-2 can succeed
        assert!(
            !successes.is_empty() || !failures.is_empty(),
            "at least some operations should complete"
        );

        // Final count should not exceed max_shards
        assert!(
            mapper.active_shard_count() <= 10,
            "should never exceed max_shards, got {}",
            mapper.active_shard_count()
        );
    }

    /// BUG #64: multiple `scale_up_provisioning(1)` calls must
    /// never push `active_count` past `max_shards` via subsequent
    /// `activate()` calls. Pre-fix the budget gate
    /// (`check_scale_up_budget`) only counted ALREADY-active
    /// shards, so several `scale_up_provisioning` calls could each
    /// pass and then each `activate()` unconditionally bumped
    /// `active_count` past the cap.
    ///
    /// Setup: at the budget edge, allocate two Provisioning shards
    /// in sequence (both pass the gate because `active_count`
    /// hasn't moved). The first `activate()` fills the budget;
    /// the second must surface `AtMaxShards` rather than overflow.
    #[test]
    fn activate_rejects_when_active_count_would_exceed_max_shards() {
        let policy = ScalingPolicy {
            min_shards: 1,
            max_shards: 4,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(3, 1024, policy).unwrap();
        // active_count = 3, max = 4. Allocate TWO Provisioning
        // shards; both pass the budget gate (it sees active_count=3).
        let ids_a = mapper.scale_up_provisioning(1).unwrap();
        let ids_b = mapper.scale_up_provisioning(1).unwrap();
        assert_eq!(ids_a.len(), 1);
        assert_eq!(ids_b.len(), 1);

        // First activate succeeds (active_count goes 3 → 4 = max).
        mapper
            .activate(ids_a[0])
            .expect("first activate must succeed");
        assert_eq!(mapper.active_shard_count(), 4);

        // Second activate must REFUSE — it would push past max.
        let err = mapper
            .activate(ids_b[0])
            .expect_err("second activate must reject (BUG #64)");
        assert!(
            matches!(err, ScalingError::AtMaxShards),
            "expected AtMaxShards, got {:?}",
            err
        );
        // active_count must still be at the cap, not over it.
        assert_eq!(mapper.active_shard_count(), 4);
    }

    /// Two concurrent `scale_up(1)` calls must never both succeed
    /// inside a single cooldown window. Before the fix, the
    /// cooldown check happened only under a read lock that was
    /// released *before* the mutating write lock, so two threads
    /// could both observe `last_scaling=None` (or stale), both
    /// acquire the write lock in turn, and both succeed — racing
    /// past the cooldown floor and (on a max-shard-bounded
    /// scenario) potentially the `max_shards` cap.
    ///
    /// Pin: across `ITERATIONS` rounds of two-thread races, every
    /// iteration sees exactly one success and one `InCooldown`.
    #[test]
    fn cooldown_is_enforced_under_concurrent_scale_up() {
        use std::sync::Arc;
        use std::sync::Barrier;
        use std::thread;

        const ITERATIONS: usize = 1_000;

        for iter in 0..ITERATIONS {
            let policy = ScalingPolicy {
                min_shards: 1,
                max_shards: 16,
                // Large cooldown so a single iteration can't pass
                // the floor between the two threads' calls.
                cooldown: Duration::from_secs(60),
                ..Default::default()
            };
            let mapper = Arc::new(ShardMapper::new(2, 1024, policy).unwrap());
            let barrier = Arc::new(Barrier::new(2));

            let m1 = mapper.clone();
            let b1 = barrier.clone();
            let m2 = mapper.clone();
            let b2 = barrier.clone();

            let h1 = thread::spawn(move || {
                b1.wait();
                m1.scale_up(1)
            });
            let h2 = thread::spawn(move || {
                b2.wait();
                m2.scale_up(1)
            });

            let r1 = h1.join().unwrap();
            let r2 = h2.join().unwrap();

            let oks = [&r1, &r2].iter().filter(|r| r.is_ok()).count();
            let cooldowns = [&r1, &r2]
                .iter()
                .filter(|r| matches!(r, Err(ScalingError::InCooldown)))
                .count();

            assert_eq!(
                oks, 1,
                "iter {iter}: expected exactly one Ok, got r1={r1:?}, r2={r2:?}"
            );
            assert_eq!(
                cooldowns, 1,
                "iter {iter}: expected exactly one InCooldown, got r1={r1:?}, r2={r2:?}"
            );
            assert_eq!(
                mapper.active_shard_count(),
                3,
                "iter {iter}: cooldown violated — both calls mutated state (shard count {})",
                mapper.active_shard_count()
            );
        }
    }

    #[test]
    fn test_scale_up_overflow_protection() {
        // Create mapper with a high starting shard ID to test overflow
        // protection. The monotonic id allocator advances by 1 per
        // shard regardless of which shards have been removed, so we
        // bump `next_shard_id` directly to simulate a near-`u16::MAX`
        // state.
        let policy = ScalingPolicy {
            max_shards: u16::MAX,
            ..Default::default()
        };
        let mapper = ShardMapper::new(1, 1024, policy).unwrap();

        // Position the allocator so the next id is u16::MAX - 1.
        // Trying to allocate 3 ids would need {MAX-1, MAX, MAX+1};
        // the last is unrepresentable in `u16`, so scale_up rejects.
        mapper
            .next_shard_id
            .store(u16::MAX - 1, AtomicOrdering::Relaxed);

        let result = mapper.scale_up(3);
        assert!(matches!(result, Err(ScalingError::AtMaxShards)));

        // Adding 1 shard should still work (id = MAX - 1).
        let result = mapper.scale_up(1);
        assert!(result.is_ok());
    }

    #[test]
    fn test_evaluate_scaling_auto_scale_disabled() {
        let policy = ScalingPolicy {
            auto_scale: false,
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::None));
    }

    #[test]
    fn test_evaluate_scaling_in_cooldown() {
        let policy = ScalingPolicy {
            cooldown: Duration::from_secs(60),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Trigger a scaling operation to start cooldown
        *mapper.last_scaling.write() = Some(Instant::now());

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::None));
    }

    #[test]
    fn test_evaluate_scaling_scale_up_on_high_fill_ratio() {
        let policy = ScalingPolicy {
            fill_ratio_threshold: 0.7,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set high fill ratio on majority of shards
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.fill_ratio = 0.9; // Above threshold
            }
        }

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::ScaleUp(_)));
    }

    #[test]
    fn test_evaluate_scaling_scale_up_on_high_latency() {
        let policy = ScalingPolicy {
            push_latency_threshold_ns: 10,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set high push latency on majority of shards
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.avg_push_latency_ns = 100; // Above threshold
            }
        }

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::ScaleUp(_)));
    }

    #[test]
    fn test_evaluate_scaling_scale_down_on_underutilized() {
        let policy = ScalingPolicy {
            underutilized_threshold: 0.2,
            min_shards: 2,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set low fill ratio and zero event rate on majority of shards
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.fill_ratio = 0.05; // Below threshold
                shard.last_metrics.event_rate = 0;
            }
        }

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::ScaleDown(_)));
    }

    #[test]
    fn test_evaluate_scaling_no_scale_up_at_max() {
        let policy = ScalingPolicy {
            fill_ratio_threshold: 0.7,
            max_shards: 4, // Already at max
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set high fill ratio
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.fill_ratio = 0.9;
            }
        }

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::None));
    }

    #[test]
    fn test_evaluate_scaling_no_scale_down_at_min() {
        let policy = ScalingPolicy {
            underutilized_threshold: 0.2,
            min_shards: 4, // Already at min
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set underutilized metrics
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.fill_ratio = 0.05;
                shard.last_metrics.event_rate = 0;
            }
        }

        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::None));
    }

    #[test]
    fn test_evaluate_scaling_ignores_draining_shards() {
        let policy = ScalingPolicy {
            fill_ratio_threshold: 0.7,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(4, 1024, policy).unwrap();

        // Set high fill ratio but mark shards as draining
        {
            let mut shards = mapper.shards.write();
            for shard in shards.iter_mut() {
                shard.last_metrics.fill_ratio = 0.9;
                shard.state = ShardState::Draining;
            }
        }

        // Should not scale up since draining shards are ignored
        let decision = mapper.evaluate_scaling();
        assert!(matches!(decision, ScalingDecision::None));
    }

    #[test]
    fn test_shard_metrics_new() {
        let metrics = ShardMetrics::new(5);
        assert_eq!(metrics.shard_id, 5);
        assert_eq!(metrics.fill_ratio, 0.0);
        assert_eq!(metrics.event_rate, 0);
        assert!(!metrics.draining);
    }

    #[test]
    fn test_shard_metrics_compute_weight() {
        let mut metrics = ShardMetrics::new(0);
        metrics.fill_ratio = 0.5;
        metrics.avg_push_latency_ns = 100;
        metrics.event_rate = 1_000_000;

        metrics.compute_weight();
        assert!(metrics.weight > 0.0);
    }

    #[test]
    fn test_shard_metrics_draining_max_weight() {
        let mut metrics = ShardMetrics::new(0);
        metrics.draining = true;
        metrics.compute_weight();
        assert_eq!(metrics.weight, f64::MAX);
    }

    #[test]
    fn test_scaling_decision_debug() {
        let none = ScalingDecision::None;
        let up = ScalingDecision::ScaleUp(2);
        let down = ScalingDecision::ScaleDown(1);

        assert!(format!("{:?}", none).contains("None"));
        assert!(format!("{:?}", up).contains("ScaleUp"));
        assert!(format!("{:?}", down).contains("ScaleDown"));
    }

    /// Regression: BUG_REPORT.md #46 — `scale_up_provisioning`
    /// allocates a shard but `select_shard` must NOT route to it
    /// until `activate` has been called. This is the load-bearing
    /// invariant that lets `EventBus::add_shard_internal` wire up
    /// drain workers before producers can land in the new ring
    /// buffer.
    #[test]
    fn provisioning_shard_is_not_selectable_until_activated() {
        let policy = ScalingPolicy {
            min_shards: 1,
            max_shards: 16,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();

        // Allocate a provisioning shard. Existing ids are 0,1; new
        // is 2.
        let new_ids = mapper.scale_up_provisioning(1).unwrap();
        assert_eq!(new_ids, vec![2]);

        // The provisioning shard must NOT appear in active accounting.
        assert_eq!(mapper.active_shard_count(), 2);
        assert_eq!(mapper.shard_state(2), Some(ShardState::Provisioning));

        // Across many hashes, `select_shard` must never pick id 2.
        let mut seen: std::collections::HashSet<u16> = std::collections::HashSet::new();
        for hash in 0u64..10_000 {
            seen.insert(mapper.select_shard(hash));
        }
        assert!(
            !seen.contains(&2),
            "provisioning shard 2 was selected — \
             producers would push into a ring buffer with no consumer (#46)"
        );
        assert!(seen == [0, 1].into_iter().collect());

        // After `activate`, the shard is selectable.
        mapper.activate(2).unwrap();
        assert_eq!(mapper.shard_state(2), Some(ShardState::Active));
        assert_eq!(mapper.active_shard_count(), 3);

        let mut seen_after: std::collections::HashSet<u16> = std::collections::HashSet::new();
        for hash in 0u64..10_000 {
            seen_after.insert(mapper.select_shard(hash));
        }
        assert!(
            seen_after.contains(&2),
            "after activate, shard 2 should be a valid select_shard target"
        );
    }

    /// Regression: BUG_REPORT.md #51 — when no `Active` shard exists
    /// (e.g. all shards are Draining or Provisioning), `select_shard`
    /// must NOT fall back to a Draining shard. Pushing into a
    /// draining shard increments `pushes_since_drain_start` and
    /// blocks finalization indefinitely.
    #[test]
    fn select_shard_does_not_fall_back_to_draining() {
        let policy = ScalingPolicy {
            min_shards: 0,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        // Force min_shards = 0 via direct construction so we can drain
        // every shard.
        let mut policy = policy;
        policy.min_shards = 1;
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();

        // Force every shard into Draining state directly. We bypass
        // `scale_down`'s min_shards floor by mutating the test-only
        // accessor; this models a state that can otherwise be
        // reached by `drain_specific` calls.
        {
            let mut shards = mapper.shards.write();
            for s in shards.iter_mut() {
                s.state = ShardState::Draining;
                s.metrics.set_draining(true);
            }
        }

        // The fallback must return the OOB sentinel `u16::MAX` rather
        // than a draining shard id (0 or 1). The upstream
        // `resolve_idx` path will see no match for `u16::MAX` and
        // surface as `Unrouted` (#44), which is the correct signal:
        // "no destination, do not push."
        for hash in 0u64..1_000 {
            let picked = mapper.select_shard(hash);
            assert_eq!(
                picked,
                u16::MAX,
                "fallback returned a draining shard ({}); pushes there would \
                 deadlock finalize_draining (#51)",
                picked
            );
        }
    }

    /// Regression: BUG_REPORT.md #32 — `scale_up(0)` previously
    /// bumped the cooldown timestamp and could spuriously fail at
    /// `u16::MAX` even though zero ids were being allocated.
    #[test]
    fn scale_up_zero_is_a_noop() {
        let policy = ScalingPolicy {
            cooldown: Duration::from_secs(60),
            ..Default::default()
        };
        let mapper = ShardMapper::new(2, 1024, policy).unwrap();
        // Pretend we just scaled — cooldown is active.
        *mapper.last_scaling.write() = Some(Instant::now());

        // scale_up(0) must not return InCooldown and must not bump
        // last_scaling.
        let before_ts = *mapper.last_scaling.read();
        let r = mapper.scale_up(0);
        assert!(
            r.is_ok(),
            "scale_up(0) should succeed as a no-op, got {r:?}"
        );
        assert_eq!(r.unwrap().len(), 0);
        let after_ts = *mapper.last_scaling.read();
        assert_eq!(before_ts, after_ts, "scale_up(0) bumped cooldown timestamp");

        // Also: position the allocator at u16::MAX and verify a
        // count==0 call doesn't trip the sentinel check.
        mapper
            .next_shard_id
            .store(u16::MAX, AtomicOrdering::Relaxed);
        assert!(mapper.scale_up(0).is_ok());
    }

    /// Regression: BUG_REPORT.md #48 — `drain_specific` must
    /// transition the shard's `MappedShard.state` to `Draining` so
    /// that `select_shard` stops routing to it. The previous
    /// `drain_shard` only flipped a metrics atomic.
    #[test]
    fn drain_specific_takes_shard_out_of_select() {
        let policy = ScalingPolicy {
            min_shards: 1,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = ShardMapper::new(3, 1024, policy).unwrap();

        // Drain shard 0.
        mapper.drain_specific(0).unwrap();
        assert_eq!(mapper.shard_state(0), Some(ShardState::Draining));
        assert_eq!(mapper.active_shard_count(), 2);

        // Across many hashes, select_shard must not pick id 0.
        for hash in 0u64..10_000 {
            let picked = mapper.select_shard(hash);
            assert_ne!(
                picked, 0,
                "select_shard returned a Draining shard id 0 — \
                 producer pushes there would block finalize_draining (#48)"
            );
        }
    }

    /// Regression: BUG_REPORT.md #33 — `set_draining(true)` and a
    /// concurrent `record_push` race on `pushes_since_drain_start`.
    /// The race itself can't be eliminated without a CAS loop on
    /// the counter (which would penalize the hot path), so the
    /// best we can pin is: after the dust settles, the counter
    /// value is bounded — never larger than the number of pushes
    /// that genuinely overlapped the transition. This catches a
    /// regression where the store-zero stops happening, where the
    /// flag publish stops happening, or where future code adds
    /// drift that compounds the race across many transitions.
    #[test]
    fn set_draining_resets_counter_under_concurrent_pushes() {
        use std::sync::Barrier;
        use std::thread;

        const ITERATIONS: usize = 200;
        const PUSHERS: usize = 4;

        for _ in 0..ITERATIONS {
            let collector = Arc::new(ShardMetricsCollector::new(0, 1024));

            // Sanity: counter starts at zero.
            assert_eq!(collector.pushes_since_drain_start(), 0);

            // Pre-load with a noticeable amount of "before drain"
            // pushes so the reset has work to do.
            for _ in 0..50 {
                collector.record_push(1);
            }
            assert_eq!(collector.pushes_since_drain_start(), 50);

            // Race: every pusher hammers `record_push` while one
            // thread calls `set_draining(true)`. After the barrier,
            // we want to observe that the counter ends up bounded
            // by PUSHERS (the number of pushes that genuinely
            // overlapped the transition) — and crucially NOT 50+,
            // which is what the buggy code with a missing reset
            // would leave behind.
            let barrier = Arc::new(Barrier::new(PUSHERS + 1));
            let mut handles = Vec::with_capacity(PUSHERS);
            for _ in 0..PUSHERS {
                let c = Arc::clone(&collector);
                let b = Arc::clone(&barrier);
                handles.push(thread::spawn(move || {
                    b.wait();
                    c.record_push(1);
                }));
            }

            barrier.wait();
            collector.set_draining(true);
            for h in handles {
                h.join().unwrap();
            }

            // After reset + at-most-PUSHERS racing pushes, the
            // counter is bounded. The strong guarantee we want is
            // "the 50 pre-drain pushes did NOT survive the reset."
            // Anything <= PUSHERS is acceptable — the race may
            // count any subset of the racing pushes.
            let final_count = collector.pushes_since_drain_start();
            assert!(
                final_count <= PUSHERS as u64,
                "set_draining reset is broken: counter is {} after reset, \
                 expected at most {} (#33)",
                final_count,
                PUSHERS
            );
        }
    }

    /// Regression: BUG_REPORT.md #49 — `finalize_draining` must
    /// drop the `shards` write lock before calling `on_shard_removed`,
    /// so a callback that re-enters the mapper (read methods like
    /// `shard_state`, `active_shard_ids`, etc.) does not deadlock.
    #[test]
    fn finalize_draining_does_not_deadlock_on_callback_reentry() {
        use std::sync::Mutex;
        let policy = ScalingPolicy {
            min_shards: 1,
            max_shards: 8,
            cooldown: Duration::from_nanos(1),
            ..Default::default()
        };
        let mapper = Arc::new(ShardMapper::new(2, 1024, policy).unwrap());

        // Set a callback that re-enters the mapper. Before the fix
        // this acquires `shards.read()` while finalize_draining
        // still holds `shards.write()`, deadlocking on parking_lot's
        // non-recursive RwLock.
        type Observation = (u16, Option<ShardState>);
        let observed_states: Arc<Mutex<Vec<Observation>>> = Arc::new(Mutex::new(Vec::new()));
        {
            let mapper_for_cb = Arc::clone(&mapper);
            let observed = Arc::clone(&observed_states);
            mapper.set_on_shard_removed(move |id| {
                let st = mapper_for_cb.shard_state(id);
                observed.lock().unwrap().push((id, st));
            });
        }

        // Drive a shard all the way to Stopped: drain it, mark its
        // metrics empty (current_len = 0), and drop drain_started
        // far enough back that the 100ms gate is satisfied.
        mapper.drain_specific(0).unwrap();
        {
            let mut shards = mapper.shards.write();
            let s = shards.iter_mut().find(|s| s.id == 0).unwrap();
            s.metrics.current_len.store(0, AtomicOrdering::Relaxed);
            s.metrics
                .pushes_since_drain_start
                .store(0, AtomicOrdering::Relaxed);
            // Backdate drain_started so the elapsed > 100ms gate trips.
            s.drain_started = Some(Instant::now() - Duration::from_secs(1));
        }

        // The call below would deadlock with the bug present.
        let stopped = mapper.finalize_draining();
        assert_eq!(stopped, vec![0]);

        // The callback must have run AND been able to read state
        // (proving the lock was released).
        let observed = observed_states.lock().unwrap().clone();
        assert_eq!(observed.len(), 1);
        assert_eq!(observed[0].0, 0);
        // State should be `Stopped` (set by finalize_draining before
        // the lock was dropped).
        assert_eq!(observed[0].1, Some(ShardState::Stopped));
    }
}
