//! Scatter-gather query routing for entity-sharded cluster deployments.
//!
//! When entities are hash-partitioned across shards, a RETRIEVE or SEARCH
//! query must fan out to all shards, collect per-shard results, and merge
//! them into a single ranked result set. This module implements:
//!
//! - **Entity-sharded write routing**: `hash(entity_id) % num_shards`
//! - **Scatter-gather RETRIEVE**: fan out to all shards, K-way merge by score
//! - **Scatter-gather SEARCH**: fan out to all shards, merge by score
//! - **Deadline propagation**: configurable budget with per-hop overhead subtracted
//! - **Partial failure**: degraded results when some shards are unreachable
//!
//! # Size / cohesion note (M0–M10 review Maintainability-S)
//!
//! This file is large (~1.5k lines), but roughly half is its in-module test
//! suite and the production code is a single tightly-coupled concern: the
//! fan-out → gather → dedup → reconcile → diversity → truncate merge pipeline,
//! whose helpers share the [`Sourced`]/[`GatherState`]/[`MergeItem`] types and
//! must stay consistent. Splitting it would fragment that pipeline without
//! reducing real complexity, so the split is deferred (not forced mid-campaign).
//! Tracked for a dedicated follow-up.

use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Condvar, Mutex, OnceLock, mpsc},
    time::{Duration, Instant},
};

use tidaldb::{
    query::{
        retrieve::{Results as RetrieveResults, Retrieve, RetrieveResult},
        search::{Search, SearchResults},
    },
    replication::shard::{RegionId, ShardId, ShardRouter},
    schema::EntityId,
    testing::SimulatedCluster,
};

use crate::error::{Result, ServerError};

/// Default query deadline (50ms as per spec Section 7.4).
const DEFAULT_DEADLINE_MS: u64 = 50;

/// Estimated network overhead per shard hop (subtracted from deadline).
const NETWORK_OVERHEAD_MS: u64 = 5;

/// Server-side ceiling on a client-supplied `deadline_ms`.
///
/// The scatter-gather budget is fully client-controlled (the `?deadline_ms=`
/// query param flows straight into [`scatter_gather_retrieve`] /
/// [`scatter_gather_search`]). An unbounded value lets a single request pin a
/// blocking-pool worker for an arbitrarily long time — a trivial resource-
/// exhaustion / slow-loris vector against the whole node. We clamp every
/// request to this ceiling (10s), which is already two orders of magnitude
/// above the 50ms spec target, so it never constrains a legitimate query while
/// capping the worst case. See [`clamp_deadline_ms`].
const MAX_DEADLINE_MS: u64 = 10_000;

/// Clamp a client-supplied scatter-gather deadline to [`MAX_DEADLINE_MS`].
///
/// `None` keeps the [`DEFAULT_DEADLINE_MS`] default. Any explicit value above
/// the ceiling is logged once and reduced, so a malicious or buggy client
/// cannot hold a worker indefinitely.
fn clamp_deadline_ms(requested: Option<u64>) -> u64 {
    match requested {
        None => DEFAULT_DEADLINE_MS,
        Some(ms) if ms > MAX_DEADLINE_MS => {
            tracing::warn!(
                requested_ms = ms,
                cap_ms = MAX_DEADLINE_MS,
                "scatter-gather deadline_ms exceeds server cap; clamping"
            );
            MAX_DEADLINE_MS
        }
        Some(ms) => ms,
    }
}

// ── Global scatter-gather worker bound ───────────────────────────────────────

/// Floor on the total concurrent shard-worker threads the process may run.
/// Even on a single-core host the fan-out gets meaningful parallelism.
const MIN_SHARD_WORKERS: usize = 8;
/// Multiplier applied to available parallelism to size the cap. Shard workers
/// are query/IO-bound (a blocking `TidalDb` read), not purely CPU-bound, so a
/// modest oversubscription keeps cores busy without unbounded growth.
const SHARD_WORKERS_PER_CORE: usize = 8;
/// Hard ceiling on the total concurrent shard-worker threads, independent of
/// core count, so a many-core host still cannot spawn an unbounded thread set
/// under a query storm.
const MAX_SHARD_WORKERS: usize = 256;

/// A process-wide counting semaphore bounding the number of scatter-gather
/// shard-worker threads that may execute a blocking shard query at once.
///
/// The router-level [`ConcurrencyLimitLayer`](tower::limit::ConcurrencyLimitLayer)
/// caps in-flight HTTP requests, but each sharded request still fans out one
/// detached thread per live shard. Without a fan-out cap, `requests × shards`
/// detached OS threads can pile up under a burst. This semaphore caps the
/// AGGREGATE concurrent shard workers regardless of how many requests fan out,
/// so the node sheds load cleanly instead of exhausting OS threads.
///
/// A worker that cannot acquire a permit before the request deadline exits
/// WITHOUT running its query; the coordinator then reports that shard as
/// degraded (never silently dropped). That is the correct behavior under
/// overload: the shard genuinely could not be serviced within budget, and the
/// would-be worker thread retires immediately rather than parking indefinitely.
struct ShardWorkerSemaphore {
    /// Available permits. Guarded by the mutex; waiters block on the condvar.
    permits: Mutex<usize>,
    available: Condvar,
}

impl ShardWorkerSemaphore {
    fn new(permits: usize) -> Self {
        Self {
            permits: Mutex::new(permits.max(1)),
            available: Condvar::new(),
        }
    }

    /// Acquire one permit, waiting at most `timeout`. Returns a guard that
    /// releases the permit on drop, or `None` if no permit became available in
    /// time (the caller should degrade rather than block further).
    fn acquire_timeout(&self, timeout: Duration) -> Option<ShardWorkerPermit<'_>> {
        let deadline = Instant::now() + timeout;
        // Poison is benign here: the only critical section is the integer
        // permit count and a notify; a panic mid-update cannot leave a torn
        // value, so recover the guard and continue (house pattern).
        let mut permits = self
            .permits
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        loop {
            if *permits > 0 {
                *permits -= 1;
                return Some(ShardWorkerPermit { sem: self });
            }
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                return None;
            }
            let (next, timed_out) = self
                .available
                .wait_timeout(permits, remaining)
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            permits = next;
            if timed_out.timed_out() && *permits == 0 {
                return None;
            }
        }
    }

    fn release(&self) {
        {
            let mut permits = self
                .permits
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            *permits += 1;
        }
        // One released permit wakes at most one waiter. Notify after dropping the
        // lock so the woken waiter does not immediately re-block on a held guard.
        self.available.notify_one();
    }
}

/// RAII permit: releases its slot back to the semaphore on drop, even if the
/// shard query panics.
struct ShardWorkerPermit<'a> {
    sem: &'a ShardWorkerSemaphore,
}

impl Drop for ShardWorkerPermit<'_> {
    fn drop(&mut self) {
        self.sem.release();
    }
}

/// The process-global shard-worker semaphore, sized once from available
/// parallelism on first use.
static SHARD_WORKER_SEMAPHORE: OnceLock<ShardWorkerSemaphore> = OnceLock::new();

/// Resolve the global shard-worker semaphore, initializing it on first use.
fn shard_worker_semaphore() -> &'static ShardWorkerSemaphore {
    SHARD_WORKER_SEMAPHORE.get_or_init(|| {
        let cores = std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get);
        let permits = cores
            .saturating_mul(SHARD_WORKERS_PER_CORE)
            .clamp(MIN_SHARD_WORKERS, MAX_SHARD_WORKERS);
        tracing::info!(permits, "scatter-gather shard-worker cap initialized");
        ShardWorkerSemaphore::new(permits)
    })
}

/// Metadata about scatter-gather query execution.
#[derive(Debug, Clone)]
pub struct ScatterGatherMeta {
    /// Whether the result set is degraded due to shard failures.
    pub degraded: bool,
    /// Shards that were unavailable during the query.
    pub unavailable_shards: Vec<String>,
    /// Number of shards that contributed results.
    pub shards_queried: usize,
    /// Total wall-clock time for the scatter-gather.
    pub elapsed_ms: u64,
    /// Per-shard deadline that was propagated.
    pub shard_deadline_ms: u64,
}

/// Determines which shard owns an entity, delegating to the engine's
/// [`ShardRouter`] so server-side write/read routing can never disagree with
/// the engine's own entity→shard mapping.
///
/// The engine [`ShardRouter::hash`] uses FNV-1a over the entity ID and returns
/// a [`ShardId`] in `0..num_shards`; we index that into `shards`, which is the
/// region/shard list in ascending order, so `ShardId(i)` maps to
/// `shards[i]`. Previously this used a divergent Knuth multiplicative hash,
/// which silently disagreed with the engine and would route the same entity to
/// different shards for writes vs. reads.
///
/// # Panics
///
/// Panics in debug mode if `shards` is empty.
#[must_use]
pub fn entity_shard(entity_id: EntityId, shards: &[RegionId]) -> RegionId {
    debug_assert!(!shards.is_empty(), "entity_shard requires non-empty shards");
    // `ShardRouter::hash` only fails on a zero shard count, which the empty
    // guard above rules out; clamp to at least one shard so the conversion is
    // infallible without an `unwrap`.
    let num_shards = u16::try_from(shards.len()).unwrap_or(u16::MAX).max(1);
    let shard =
        ShardRouter::hash(num_shards).map_or(ShardId::SINGLE, |router| router.route(entity_id));
    // ShardId(i) corresponds to shards[i]; the router guarantees i < num_shards
    // for a non-empty list. Fall back to the first shard rather than indexing
    // out of bounds if a caller ever violates the non-empty precondition.
    shards
        .get(shard.0 as usize)
        .or_else(|| shards.first())
        .copied()
        .unwrap_or(RegionId::SINGLE)
}

/// Write an item to the owning shard only (entity-sharded mode).
///
/// # Errors
///
/// Returns [`ServerError`] if the owning shard's write fails.
// `metadata` is always the std-hasher `HashMap` built by the cluster sim; a
// generic hasher bound would add noise with no caller benefit.
#[allow(clippy::implicit_hasher)]
pub fn sharded_write_item(
    cluster: &SimulatedCluster,
    entity_id: EntityId,
    metadata: &HashMap<String, String>,
    shards: &[RegionId],
) -> Result<()> {
    let shard = entity_shard(entity_id, shards);
    cluster
        .node(shard)
        .db
        .write_item_with_metadata(entity_id, metadata)
        .map_err(ServerError::from)
}

/// Write an embedding to the owning shard only (entity-sharded mode).
///
/// # Errors
///
/// Returns [`ServerError`] if the owning shard's embedding write fails.
pub fn sharded_write_embedding(
    cluster: &SimulatedCluster,
    entity_id: EntityId,
    embedding: &[f32],
    shards: &[RegionId],
) -> Result<()> {
    let shard = entity_shard(entity_id, shards);
    cluster
        .node(shard)
        .db
        .write_item_embedding(entity_id, embedding)
        .map_err(ServerError::from)
}

/// Write a signal to the owning shard (entity-sharded mode).
///
/// # Errors
///
/// Returns [`ServerError`] if the owning shard's signal write fails.
pub fn sharded_write_signal(
    cluster: &SimulatedCluster,
    signal_name: &str,
    entity_id: EntityId,
    weight: f64,
    shards: &[RegionId],
) -> Result<()> {
    let shard = entity_shard(entity_id, shards);
    cluster
        .node(shard)
        .db
        .signal(
            signal_name,
            entity_id,
            weight,
            tidaldb::schema::Timestamp::now(),
        )
        .map_err(ServerError::from)
}

/// Resolve a human-readable shard name, falling back to `s<id>`.
fn shard_name(region_names: &HashMap<RegionId, String>, shard: RegionId) -> String {
    region_names
        .get(&shard)
        .cloned()
        .unwrap_or_else(|| format!("s{}", shard.0))
}

/// What one shard contributed to a scatter-gather.
struct ShardOutcome<T> {
    items: Vec<T>,
    total_candidates: usize,
}

/// A merged item tagged with the shard ([`RegionId`]) that actually returned
/// it.
///
/// Coordinator-level diversity resolves each item's creator from `region` —
/// the node that owns the item in an entity-sharded topology — rather than from
/// a single fixed replica. In a replicated topology every shard holds every
/// entity, so `region` is simply one healthy replica that has the item; in an
/// entity-sharded topology it is the *only* node that has it. Either way the
/// creator lookup hits a node that actually stores the item. See
/// [`enforce_max_per_creator`].
struct Sourced<T> {
    /// The shard that returned this item (and therefore can resolve its
    /// metadata locally).
    region: RegionId,
    item: T,
}

/// Per-shard query state accumulated by [`dispatch_shards`].
struct GatherState<T> {
    /// Each merged item paired with the shard that returned it, so the merge
    /// stage can resolve creators from the owning node (entity-sharded) or a
    /// holding replica (replicated). See [`Sourced`].
    items: Vec<Sourced<T>>,
    /// Each contributing shard's `total_candidates`, kept separately so the
    /// merge stage can reconcile the count instead of blindly summing — a sum
    /// double-counts the candidate universe across REPLICATED shards (every
    /// shard sees every entity). See [`reconcile_total_candidates`].
    per_shard_totals: Vec<usize>,
    unavailable_shards: Vec<String>,
    shards_queried: usize,
}

/// Fan out a per-shard query CONCURRENTLY and gather the results under a hard
/// total-time budget.
///
/// Each non-partitioned shard runs `query_one` on its own **detached** OS
/// thread holding an owned `Arc<SimulatedCluster>` clone. The underlying
/// `TidalDb` query is a *blocking* call, so true thread-level concurrency (not
/// cooperative async) is required for one slow shard not to serialize behind
/// the others — and the threads must be detached, not scoped, so the
/// coordinator can return its partial result the instant the budget expires
/// without joining a shard whose blocking query is still in flight. (A scoped
/// join would re-introduce the very hang this fix removes.) The coordinator
/// drains the result channel with [`mpsc::Receiver::recv_timeout`] against the
/// remaining budget.
///
/// Shards that error, or that fail to report by the deadline, are recorded in
/// `unavailable_shards` (degraded) — they are NEVER silently dropped. A
/// detached worker that finishes after the deadline simply sends into a
/// receiver that has already been dropped; the send fails harmlessly and the
/// worker exits. Because every worker only *reads* the shared cluster, leaking
/// a still-running worker past the request is sound.
///
/// # Fan-out cap
///
/// Each worker must acquire a permit from the process-global
/// [`shard_worker_semaphore`] before running its blocking query, so the
/// AGGREGATE number of concurrently-executing shard workers is bounded
/// regardless of how many sharded requests fan out at once. A worker that
/// cannot get a permit within the remaining budget retires immediately and is
/// reported degraded — never silently dropped, and never left parked
/// indefinitely. Together with the router's request-concurrency limit this caps
/// total detached threads under a query storm instead of growing them without
/// bound.
fn dispatch_shards<T, F>(
    cluster: &Arc<SimulatedCluster>,
    shards: &[RegionId],
    region_names: &HashMap<RegionId, String>,
    deadline: Duration,
    query_one: F,
) -> GatherState<T>
where
    T: Send + 'static,
    F: Fn(&SimulatedCluster, RegionId) -> Result<ShardOutcome<T>> + Clone + Send + 'static,
{
    let start = Instant::now();
    let mut state = GatherState {
        items: Vec::new(),
        per_shard_totals: Vec::new(),
        unavailable_shards: Vec::new(),
        shards_queried: 0,
    };

    // Pre-filter partitioned shards: they are known-unavailable, so we never
    // spawn a worker for them.
    let live_shards: Vec<RegionId> = shards
        .iter()
        .copied()
        .filter(|&shard| {
            if cluster.is_partitioned(shard) {
                state
                    .unavailable_shards
                    .push(shard_name(region_names, shard));
                false
            } else {
                true
            }
        })
        .collect();

    if live_shards.is_empty() {
        return state;
    }

    // `(shard, Result<ShardOutcome>)` flows back over this channel. The bound
    // equals the live-shard count so no worker ever blocks on send — and so a
    // late worker can always deposit its result without waiting on a receiver
    // that may already be gone.
    let (tx, rx) = mpsc::sync_channel::<(RegionId, Result<ShardOutcome<T>>)>(live_shards.len());

    // Shards whose worker actually started. A shard whose thread fails to spawn
    // is marked degraded immediately and excluded from the wait set so we never
    // burn the whole deadline waiting on a result that can never arrive.
    let mut dispatched: Vec<RegionId> = Vec::with_capacity(live_shards.len());
    let sem = shard_worker_semaphore();
    for &shard in &live_shards {
        let tx = tx.clone();
        let cluster = Arc::clone(cluster);
        let query_one = query_one.clone();
        let spawned = std::thread::Builder::new()
            .name(format!("scatter-shard-{}", shard.0))
            .spawn(move || {
                // Bound the AGGREGATE concurrent shard workers process-wide: a
                // worker that cannot get a permit before the budget expires
                // retires immediately and reports degraded, rather than parking
                // a thread or running an out-of-budget query. The permit is held
                // for exactly the blocking query and released on drop (even on
                // panic). See [`shard_worker_semaphore`].
                let remaining = deadline.saturating_sub(start.elapsed());
                // `_permit` (RAII) is held for exactly the blocking query and
                // released when the closure returns; `None` means the cap was hit
                // before the deadline, so this worker retires as degraded.
                let outcome = sem.acquire_timeout(remaining).map_or_else(
                    || {
                        Err(ServerError::Unavailable(
                            "scatter-gather worker cap reached before deadline".into(),
                        ))
                    },
                    |_permit| query_one(&cluster, shard),
                );
                // The receiver may already have moved on after the deadline; a
                // closed channel is expected and benign, so the error is dropped.
                let _ = tx.send((shard, outcome));
            });
        match spawned {
            Ok(_handle) => dispatched.push(shard),
            Err(e) => {
                // OS thread exhaustion: degrade this shard rather than hang or
                // panic. The other shards still race normally.
                let name = shard_name(region_names, shard);
                tracing::error!(shard = shard.0, region = %name, error = %e, "failed to spawn scatter-gather worker; marking degraded");
                state.unavailable_shards.push(name);
            }
        }
    }
    // Drop the coordinator's own sender so the channel closes once every worker
    // has sent (and been dropped), letting `recv_timeout` observe `Disconnected`
    // instead of waiting out the full budget when all shards have reported.
    drop(tx);

    // Shards whose result we have folded in (success OR error). Anything in
    // `dispatched` but not here when the budget expires is a timed-out shard.
    let mut reported: HashSet<RegionId> = HashSet::with_capacity(dispatched.len());

    while reported.len() < dispatched.len() {
        let remaining = deadline.saturating_sub(start.elapsed());
        if remaining.is_zero() {
            break;
        }
        match rx.recv_timeout(remaining) {
            Ok((shard, outcome)) => {
                fold_outcome(&mut state, &mut reported, region_names, shard, outcome);
            }
            Err(mpsc::RecvTimeoutError::Timeout | mpsc::RecvTimeoutError::Disconnected) => {
                break;
            }
        }
    }

    // The budget elapsed (or the loop broke). Drain anything that landed in the
    // race window so a straggler that finished just in time is not mis-reported
    // as a timeout. We do NOT block here — the receiver is dropped right after,
    // so any still-running worker's later send fails harmlessly.
    while reported.len() < dispatched.len() {
        match rx.try_recv() {
            Ok((shard, outcome)) => {
                fold_outcome(&mut state, &mut reported, region_names, shard, outcome);
            }
            Err(_) => break,
        }
    }

    // Any dispatched shard with no folded result missed the deadline. Mark it
    // degraded by name — never silently truncate.
    let timed_out: Vec<RegionId> = dispatched
        .iter()
        .copied()
        .filter(|shard| !reported.contains(shard))
        .collect();
    if !timed_out.is_empty() {
        tracing::warn!(
            dropped = timed_out.len(),
            "scatter-gather deadline exceeded; {} shard(s) dropped as degraded",
            timed_out.len()
        );
        for shard in timed_out {
            state
                .unavailable_shards
                .push(shard_name(region_names, shard));
        }
    }

    state
}

/// Fold one shard's reported outcome into the gather state and mark it reported.
fn fold_outcome<T>(
    state: &mut GatherState<T>,
    reported: &mut HashSet<RegionId>,
    region_names: &HashMap<RegionId, String>,
    shard: RegionId,
    outcome: Result<ShardOutcome<T>>,
) {
    reported.insert(shard);
    match outcome {
        Ok(outcome) => {
            state.per_shard_totals.push(outcome.total_candidates);
            // Tag every item with the shard that returned it so coordinator-
            // level diversity can resolve its creator from a node that actually
            // stores the item (the owning shard in entity-sharded mode).
            state
                .items
                .extend(outcome.items.into_iter().map(|item| Sourced {
                    region: shard,
                    item,
                }));
            state.shards_queried += 1;
        }
        Err(e) => {
            let name = shard_name(region_names, shard);
            tracing::warn!(shard = shard.0, region = %name, error = %e, "shard query failed; marking degraded");
            state.unavailable_shards.push(name);
        }
    }
}

// ── Merge helpers: dedup + candidate-count reconciliation ────────────────────

/// A merged scatter-gather item that exposes the identity, score, and rank slot
/// the coordinator needs to dedup replicated copies, merge by score, and assign
/// 1-based ranks after the merge.
trait MergeItem {
    /// The entity this result is for. Replicated shards return the SAME entity,
    /// so the coordinator dedups on this.
    fn entity_id(&self) -> EntityId;
    /// The (already-normalized) score used for K-way merge ordering.
    fn score(&self) -> f64;
    /// Assign the post-merge 1-based rank, so the shared merge tail can re-rank
    /// either result type without duplicating the loop per query kind.
    fn set_rank(&mut self, rank: usize);
}

impl MergeItem for RetrieveResult {
    fn entity_id(&self) -> EntityId {
        self.entity_id
    }
    fn score(&self) -> f64 {
        self.score
    }
    fn set_rank(&mut self, rank: usize) {
        self.rank = rank;
    }
}

impl MergeItem for tidaldb::query::search::SearchResultItem {
    fn entity_id(&self) -> EntityId {
        self.entity_id
    }
    fn score(&self) -> f64 {
        self.score
    }
    fn set_rank(&mut self, rank: usize) {
        self.rank = rank;
    }
}

/// Collapse duplicate entities returned by REPLICATED shards, keeping the
/// highest-scoring copy of each entity.
///
/// In a replicated topology every shard holds every entity, so the same entity
/// can appear once per shard. Without this the merged result would list the
/// same item up to `num_shards` times. The surviving copy carries its own
/// source [`RegionId`] (the shard that returned the best-scoring copy), which a
/// later creator lookup uses — that node demonstrably holds the entity, so the
/// metadata read can never miss. Returns the deduped item vector and `true` if
/// any duplicate was collapsed (i.e. shards overlapped) — the caller uses that
/// signal to reconcile `total_candidates`.
fn dedup_by_entity<T: MergeItem>(items: Vec<Sourced<T>>) -> (Vec<Sourced<T>>, bool) {
    // entity_id → index of the best-scoring copy seen so far in `out`.
    let mut best: HashMap<u64, usize> = HashMap::with_capacity(items.len());
    let mut out: Vec<Sourced<T>> = Vec::with_capacity(items.len());
    let mut overlap = false;
    for sourced in items {
        let key = sourced.item.entity_id().as_u64();
        if let Some(idx) = best.get(&key).copied() {
            overlap = true;
            if sourced.item.score() > out[idx].item.score() {
                out[idx] = sourced;
            }
        } else {
            best.insert(key, out.len());
            out.push(sourced);
        }
    }
    (out, overlap)
}

/// Reconcile the merged `total_candidates` so REPLICATED shards are not
/// counted multiple times.
///
/// - **Replicated** (`overlap_detected`, i.e. [`dedup_by_entity`] collapsed at
///   least one entity returned by two shards): every shard considered the same
///   candidate universe, so the distinct count is the LARGEST single-shard
///   total, not the sum — summing inflated it by up to `num_shards`x.
/// - **Entity-sharded** (no overlap across the returned items): each shard owns
///   a disjoint slice of entities, so the totals genuinely add up.
///
/// `deduped_len` is the post-dedup merged item count, used as a floor so the
/// reported total can never be smaller than the items actually returned.
fn reconcile_total_candidates(
    per_shard_totals: &[usize],
    deduped_len: usize,
    overlap_detected: bool,
) -> usize {
    let reconciled = if overlap_detected {
        // Replicated: the candidate universe is one replica's worth.
        per_shard_totals.iter().copied().max().unwrap_or(0)
    } else {
        // Disjoint shards: the universes add up.
        per_shard_totals.iter().sum()
    };
    reconciled.max(deduped_len)
}

/// Enforce `max_per_creator` ACROSS the merged shard results.
///
/// Each shard enforces diversity locally, but the coordinator merges several
/// shards' top-K lists, so a creator can re-appear above the cap in the merged
/// set (replicated shards each contribute the same creator; entity-sharded
/// shards each contribute their slice of a prolific creator). This walks the
/// already-score-sorted `items` and drops any item whose creator has already
/// hit the cap, mirroring the engine's per-shard `DiversitySelector` but at the
/// coordinator level.
///
/// Creators are resolved from the `creator_id` metadata of **the shard that
/// actually returned each item** ([`Sourced::region`]), not from one fixed
/// replica. This is the fix for the entity-sharded under-enforcement bug: when
/// items live on disjoint shards, the leader does not hold the ones owned by
/// other shards, so a leader-only lookup resolved their creator to `None` and
/// silently treated a prolific creator's items as uncapped. Reading from the
/// returning shard guarantees the metadata read hits a node that stores the
/// item, in both replicated and entity-sharded topologies. Items whose creator
/// still cannot be resolved (no `creator_id`, or a metadata read error) are
/// treated as having a unique, uncapped creator — never dropped — so a genuine
/// metadata gap can only ever UNDER-enforce, never hide a result.
///
/// Returns the number of items dropped, so the caller can flag the result as
/// not fully constraint-satisfied.
fn enforce_max_per_creator<T: MergeItem>(
    cluster: &SimulatedCluster,
    items: &mut Vec<Sourced<T>>,
    max_per_creator: usize,
) -> usize {
    if max_per_creator == 0 {
        // A zero cap would drop everything; treat as "no cap" rather than
        // silently emptying the feed — matches the engine, which ignores a 0.
        return 0;
    }

    let mut per_creator: HashMap<u64, usize> = HashMap::new();
    let mut dropped = 0usize;

    items.retain(|sourced| {
        // Resolve from the node that returned this item, so an entity-sharded
        // owner's metadata is always reachable.
        let creator = cluster
            .node(sourced.region)
            .db
            .get_item_metadata(sourced.item.entity_id())
            .ok()
            .flatten()
            .and_then(|meta| meta.get("creator_id").and_then(|c| c.parse::<u64>().ok()));

        // Unattributed items (None) have no creator to over-represent: keep them.
        creator.is_none_or(|cid| {
            let count = per_creator.entry(cid).or_insert(0);
            if *count >= max_per_creator {
                dropped += 1;
                false
            } else {
                *count += 1;
                true
            }
        })
    });

    dropped
}

/// The type-agnostic result of merging gathered shard items: the final ranked
/// page plus the two derived facts the caller folds into its result struct.
struct MergedResults<T> {
    /// The deduped, diversity-enforced, score-sorted, truncated, re-ranked page.
    items: Vec<T>,
    /// Candidate universe reconciled across shards (see
    /// [`reconcile_total_candidates`]) — never the raw per-shard sum.
    total_candidates: usize,
    /// `false` iff the coordinator-level diversity pass dropped any item.
    constraints_satisfied: bool,
}

/// Shared merge + assemble tail for RETRIEVE and SEARCH.
///
/// `scatter_gather_retrieve` and `scatter_gather_search` differ only in their
/// per-shard query closure and the concrete result struct they build; the merge
/// pipeline between is identical, so it lives here once:
///
/// 1. **Dedup** replicated copies of the same entity (keep the best-scoring
///    copy), recording whether shards overlapped.
/// 2. **Sort by score descending**, with `f64::total_cmp` so a NaN score yields
///    a *total*, stable order instead of being treated as equal-to-everything
///    (`partial_cmp` returns `None` for NaN, which silently degrades the sort to
///    an unstable partial order). `total_cmp` ranks NaN deterministically (a
///    positive NaN is the greatest value, so it sorts to the front of this
///    descending compare), so a single poisoned score can never interleave with
///    or scramble the ordering of the real scores.
/// 3. **Re-enforce `max_per_creator`** across the merged set when requested,
///    resolving each item's creator from the shard that returned it.
/// 4. **Reconcile `total_candidates`** BEFORE truncation so it reflects the
///    merged universe, not the page.
/// 5. **Truncate** to `limit`, drop the source-region tags, and assign 1-based
///    ranks.
fn merge_and_assemble<T: MergeItem>(
    cluster: &SimulatedCluster,
    items: Vec<Sourced<T>>,
    per_shard_totals: &[usize],
    max_per_creator: Option<usize>,
    limit: usize,
) -> MergedResults<T> {
    // Dedup replicated copies of the same entity (keep the best-scoring copy).
    let (mut items, overlap_detected) = dedup_by_entity(items);

    // Sort the deduped set by score descending. `total_cmp` gives a total order
    // even if a score is NaN, so one poisoned score cannot scramble the rest.
    items.sort_by(|a, b| b.item.score().total_cmp(&a.item.score()));

    // Re-enforce max-per-creator across the merged set (each shard only enforced
    // over its own slice). Each item's creator is resolved from the shard that
    // returned it, so an entity-sharded owner's metadata is reachable.
    let mut constraints_satisfied = true;
    if let Some(max_per_creator) = max_per_creator {
        let dropped = enforce_max_per_creator(cluster, &mut items, max_per_creator);
        if dropped > 0 {
            constraints_satisfied = false;
        }
    }

    // Reconcile the candidate count BEFORE truncating to the page limit so it
    // reflects the merged universe, not the page.
    let total_candidates =
        reconcile_total_candidates(per_shard_totals, items.len(), overlap_detected);

    // Take top limit, drop the source-region tags, and re-rank (assign 1-based
    // ranks after merge).
    let limit = limit.min(items.len());
    items.truncate(limit);
    let mut items: Vec<T> = items.into_iter().map(|s| s.item).collect();
    for (i, item) in items.iter_mut().enumerate() {
        item.set_rank(i + 1);
    }

    MergedResults {
        items,
        total_candidates,
        constraints_satisfied,
    }
}

/// Scatter-gather RETRIEVE across all shards.
///
/// Fans out the query to each non-partitioned shard CONCURRENTLY, wrapping the
/// whole gather in a hard total-time budget (per-shard deadline + network
/// overhead). The merge then, in order:
///
/// 1. **Dedups** replicated copies of the same entity (every replica returns
///    the same entity), keeping the highest-scoring copy.
/// 2. **Reconciles `total_candidates`** so replicated shards are not counted
///    multiple times (see [`reconcile_total_candidates`]) — a plain sum
///    over-counted the candidate universe by up to `num_shards`x.
/// 3. **Re-enforces `max_per_creator`** across the merged set when the query
///    declares it (see [`enforce_max_per_creator`]) — each shard only enforces
///    diversity over its own slice, and each item's creator is resolved from
///    the shard that returned it so entity-sharded owners are reachable.
/// 4. Sorts by score descending and takes the top-K.
///
/// Shards that error or miss the deadline are reported as degraded — never
/// silently truncated.
///
/// Returns the merged results and execution metadata.
///
/// # Errors
///
/// Returns [`ServerError`] if building the per-shard query fails. Shard-level
/// failures and timeouts are reported as degraded in the metadata, not as `Err`.
// `region_names` is always the std-hasher `HashMap` owned by the cluster sim.
#[allow(clippy::implicit_hasher)]
pub fn scatter_gather_retrieve(
    cluster: &Arc<SimulatedCluster>,
    query: &Retrieve,
    shards: &[RegionId],
    region_names: &HashMap<RegionId, String>,
    deadline_ms: Option<u64>,
) -> Result<(RetrieveResults, ScatterGatherMeta)> {
    let start = Instant::now();
    let budget_ms = clamp_deadline_ms(deadline_ms);
    let total_deadline = Duration::from_millis(budget_ms);
    let shard_deadline_ms = budget_ms.saturating_sub(NETWORK_OVERHEAD_MS);

    // Each detached worker holds an owned clone of the query, so it can outlive
    // this stack frame if the shard's blocking read exceeds the deadline.
    let shared_query = Arc::new(query.clone());
    let GatherState {
        items: all_items,
        per_shard_totals,
        unavailable_shards,
        shards_queried,
    } = dispatch_shards::<RetrieveResult, _>(
        cluster,
        shards,
        region_names,
        total_deadline,
        move |cluster, shard| {
            let result = cluster
                .retrieve(shard, &shared_query)
                .map_err(ServerError::from)?;
            Ok(ShardOutcome {
                items: result.items,
                total_candidates: result.total_candidates,
            })
        },
    );

    // Dedup → NaN-safe score sort → coordinator diversity → reconcile → top-K
    // re-rank: the tail shared with SEARCH (see [`merge_and_assemble`]).
    let max_per_creator = query.diversity.as_ref().and_then(|d| d.max_per_creator);
    let MergedResults {
        items: all_items,
        total_candidates,
        constraints_satisfied,
    } = merge_and_assemble(
        cluster,
        all_items,
        &per_shard_totals,
        max_per_creator,
        query.limit,
    );
    let limit = all_items.len();

    let elapsed = start.elapsed();
    let meta = ScatterGatherMeta {
        degraded: !unavailable_shards.is_empty(),
        unavailable_shards,
        shards_queried,
        elapsed_ms: elapsed.as_millis() as u64,
        shard_deadline_ms,
    };

    let results = RetrieveResults {
        items: all_items,
        next_cursor: None,
        total_candidates,
        constraints_satisfied,
        warnings: Vec::new(),
        session_snapshot: None,
        degradation_level: tidaldb::load::DegradationLevel::Full,
        stats: tidaldb::query::stats::QueryStats {
            candidates_considered: total_candidates,
            candidates_after_filter: total_candidates,
            candidates_after_diversity: limit,
            filters_applied: 0,
            scoring_time_us: 0,
            diversity_time_us: 0,
            total_time_us: elapsed.as_micros() as u64,
            degradation_level: 0,
            profile_name: String::new(),
            membership_epoch: None,
        },
        policy_metadata: tidaldb::query::retrieve::types::PolicyMetadata::default(),
    };

    Ok((results, meta))
}

/// Scatter-gather SEARCH across all shards.
///
/// Fans out the search query to each non-partitioned shard CONCURRENTLY under a
/// hard total-time budget. Each shard reloads its text index inside its own
/// worker before searching. The merge then dedups replicated copies of the
/// same entity, re-enforces `max_per_creator` across the merged set when the
/// query declares it, reconciles `total_candidates` so replicated shards are
/// not counted multiple times, then sorts by score descending and takes the
/// top-K. Shards that error or miss the deadline are reported as degraded —
/// never silently truncated.
///
/// # Errors
///
/// Returns [`ServerError`] if building the per-shard query fails. Shard-level
/// failures and timeouts are reported as degraded in the metadata, not as `Err`.
// `region_names` is always the std-hasher `HashMap` owned by the cluster sim.
#[allow(clippy::implicit_hasher)]
pub fn scatter_gather_search(
    cluster: &Arc<SimulatedCluster>,
    query: &Search,
    shards: &[RegionId],
    region_names: &HashMap<RegionId, String>,
    deadline_ms: Option<u64>,
) -> Result<(SearchResults, ScatterGatherMeta)> {
    let start = Instant::now();
    let budget_ms = clamp_deadline_ms(deadline_ms);
    let total_deadline = Duration::from_millis(budget_ms);
    let shard_deadline_ms = budget_ms.saturating_sub(NETWORK_OVERHEAD_MS);

    // Each detached worker holds an owned clone of the query, so it can outlive
    // this stack frame if the shard's blocking search exceeds the deadline.
    let shared_query = Arc::new(query.clone());
    let GatherState {
        items: all_items,
        per_shard_totals,
        unavailable_shards,
        shards_queried,
    } = dispatch_shards::<tidaldb::query::search::SearchResultItem, _>(
        cluster,
        shards,
        region_names,
        total_deadline,
        move |cluster, shard| {
            // Reload text index before searching this shard.
            if let Err(e) = cluster.node(shard).db.reload_text_index() {
                tracing::warn!(shard = shard.0, error = %e, "failed to reload text index");
            }
            let result = cluster
                .search(shard, &shared_query)
                .map_err(ServerError::from)?;
            Ok(ShardOutcome {
                items: result.items,
                total_candidates: result.total_candidates,
            })
        },
    );

    // Dedup → NaN-safe score sort → coordinator diversity → reconcile → top-K
    // re-rank: the tail shared with RETRIEVE (see [`merge_and_assemble`]).
    let max_per_creator = query.diversity.as_ref().and_then(|d| d.max_per_creator);
    let MergedResults {
        items: all_items,
        total_candidates,
        constraints_satisfied,
    } = merge_and_assemble(
        cluster,
        all_items,
        &per_shard_totals,
        max_per_creator,
        query.limit as usize,
    );
    // Preserve the prior SEARCH stat: `candidates_after_diversity` reports the
    // REQUESTED page size, not the (possibly smaller) returned count.
    let limit = query.limit as usize;

    let elapsed = start.elapsed();
    let meta = ScatterGatherMeta {
        degraded: !unavailable_shards.is_empty(),
        unavailable_shards,
        shards_queried,
        elapsed_ms: elapsed.as_millis() as u64,
        shard_deadline_ms,
    };

    let results = SearchResults {
        items: all_items,
        next_cursor: None,
        total_candidates,
        constraints_satisfied,
        warnings: Vec::new(),
        session_snapshot: None,
        degradation_level: tidaldb::load::DegradationLevel::Full,
        stats: tidaldb::query::stats::QueryStats {
            candidates_considered: total_candidates,
            candidates_after_filter: total_candidates,
            candidates_after_diversity: limit,
            filters_applied: 0,
            scoring_time_us: 0,
            diversity_time_us: 0,
            total_time_us: elapsed.as_micros() as u64,
            degradation_level: 0,
            profile_name: String::new(),
            membership_epoch: None,
        },
    };

    Ok((results, meta))
}

#[cfg(test)]
// Test exemptions: unwrap on known-good fixtures + loop-counter casts in
// distribution / sharding math are idiomatic here.
#[allow(
    clippy::unwrap_used,
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation
)]
mod tests {
    use std::time::Duration;

    use tidaldb::{
        replication::shard::RegionId,
        schema::{DecaySpec, EntityKind, SchemaBuilder, Window},
        testing::cluster::ClusterConfig,
    };

    use super::*;

    fn test_schema() -> tidaldb::schema::Schema {
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal(
                "view",
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(7 * 24 * 3600),
                },
            )
            .windows(&[Window::OneHour])
            .velocity(false)
            .add();
        builder.build().unwrap()
    }

    fn four_region_cluster() -> (
        Arc<SimulatedCluster>,
        Vec<RegionId>,
        HashMap<RegionId, String>,
    ) {
        let regions = vec![RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
        let config = ClusterConfig {
            regions: regions.clone(),
            leader_region: RegionId(0),
            schema: test_schema(),
            profiles: Vec::new(),
            transports: None,
        };
        let cluster = Arc::new(SimulatedCluster::build(config));
        let names: HashMap<RegionId, String> = regions
            .iter()
            .map(|&r| (r, format!("region-{}", r.0)))
            .collect();
        (cluster, regions, names)
    }

    #[test]
    fn entity_shard_distributes_evenly() {
        let shards = vec![RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
        let mut counts = [0u32; 4];
        for i in 0..4000u64 {
            let shard = entity_shard(EntityId::new(i), &shards);
            counts[shard.0 as usize] += 1;
        }
        // Each shard should get roughly 1000 of 4000 IDs (±20%).
        for (idx, &c) in counts.iter().enumerate() {
            assert!(c > 800 && c < 1200, "shard {idx} got {c}, expected ~1000");
        }
    }

    #[test]
    fn entity_shard_deterministic() {
        let shards = vec![RegionId(0), RegionId(1), RegionId(2)];
        let a = entity_shard(EntityId::new(42), &shards);
        let b = entity_shard(EntityId::new(42), &shards);
        assert_eq!(a, b, "same entity_id must always map to same shard");
    }

    /// AC1: RETRIEVE across 4 shards returns correct top-K merged by score.
    #[test]
    fn scatter_gather_retrieve_merges_across_shards() {
        let (cluster, shards, names) = four_region_cluster();

        // Write items with signals to each shard (replicated cluster, so all
        // shards see all data — write to leader, it replicates).
        for i in 1..=20u64 {
            let eid = EntityId::new(i);
            cluster
                .write_item_with_metadata(eid, &HashMap::new())
                .unwrap();
            cluster.write_signal("view", eid, i as f64).unwrap();
        }

        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(10)
            .build()
            .unwrap();

        let (result, meta) = scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, None)
            .expect("scatter-gather should succeed");

        // Should return items (up to 10).
        assert!(!result.items.is_empty(), "should return items");
        assert!(result.items.len() <= 10, "should respect limit");

        // Scores should be in descending order.
        for w in result.items.windows(2) {
            assert!(
                w[0].score >= w[1].score,
                "scores not descending: {} < {}",
                w[0].score,
                w[1].score
            );
        }

        // Ranks should be 1-based sequential.
        for (i, item) in result.items.iter().enumerate() {
            assert_eq!(item.rank, i + 1, "rank mismatch at position {i}");
        }

        // No degradation.
        assert!(!meta.degraded);
        assert!(meta.unavailable_shards.is_empty());
        assert_eq!(meta.shards_queried, 4);
    }

    /// AC3: One unreachable shard returns partial results with degraded=true.
    #[test]
    fn scatter_gather_degraded_when_shard_partitioned() {
        let (cluster, shards, names) = four_region_cluster();

        // Write data.
        for i in 1..=10u64 {
            let eid = EntityId::new(i);
            cluster
                .write_item_with_metadata(eid, &HashMap::new())
                .unwrap();
            cluster.write_signal("view", eid, 1.0).unwrap();
        }

        // Partition region 2.
        cluster.partition_region(RegionId(2));

        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(10)
            .build()
            .unwrap();

        let (result, meta) = scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, None)
            .expect("scatter-gather should succeed even with partitioned shard");

        // Should still return results (from 3 healthy shards).
        assert!(meta.degraded, "should be degraded");
        assert_eq!(meta.unavailable_shards.len(), 1);
        assert_eq!(meta.unavailable_shards[0], "region-2");
        assert_eq!(meta.shards_queried, 3);

        // Result should still contain items (not an error).
        // In replicated topology, all healthy shards have all data.
        assert!(!result.items.is_empty());
    }

    /// AC4: Deadline propagation subtracts network overhead.
    #[test]
    fn scatter_gather_deadline_propagation() {
        let (cluster, shards, names) = four_region_cluster();

        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(5)
            .build()
            .unwrap();

        // With default deadline (50ms).
        let (_result, meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, None).unwrap();
        assert_eq!(
            meta.shard_deadline_ms,
            DEFAULT_DEADLINE_MS - NETWORK_OVERHEAD_MS,
            "shard deadline should be total - overhead"
        );

        // With custom deadline (100ms).
        let (_result, meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(100)).unwrap();
        assert_eq!(meta.shard_deadline_ms, 100 - NETWORK_OVERHEAD_MS);

        // With very small deadline (less than overhead).
        let (_result, meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(3)).unwrap();
        assert_eq!(meta.shard_deadline_ms, 0, "should saturate at 0");
    }

    /// A client-supplied `deadline_ms` is clamped to the server-side ceiling so
    /// a single request cannot pin a worker indefinitely.
    #[test]
    fn deadline_ms_is_clamped_to_server_cap() {
        // None → default budget.
        assert_eq!(clamp_deadline_ms(None), DEFAULT_DEADLINE_MS);
        // Below the cap → passed through unchanged.
        assert_eq!(clamp_deadline_ms(Some(250)), 250);
        // Exactly at the cap → unchanged.
        assert_eq!(clamp_deadline_ms(Some(MAX_DEADLINE_MS)), MAX_DEADLINE_MS);
        // Above the cap → clamped down.
        assert_eq!(
            clamp_deadline_ms(Some(MAX_DEADLINE_MS + 1)),
            MAX_DEADLINE_MS
        );
        assert_eq!(clamp_deadline_ms(Some(u64::MAX)), MAX_DEADLINE_MS);
    }

    /// An over-budget `deadline_ms` flowing through the full retrieve path is
    /// clamped: the reported `shard_deadline_ms` reflects the capped budget, not
    /// the (absurd) requested one.
    #[test]
    fn scatter_gather_retrieve_clamps_oversized_deadline() {
        let (cluster, shards, names) = four_region_cluster();
        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(5)
            .build()
            .unwrap();
        let (_result, meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(u64::MAX)).unwrap();
        assert_eq!(
            meta.shard_deadline_ms,
            MAX_DEADLINE_MS - NETWORK_OVERHEAD_MS,
            "oversized deadline must be clamped before overhead subtraction"
        );
    }

    /// Scatter-gather search works across shards.
    #[test]
    fn scatter_gather_search_merges_across_shards() {
        let (cluster, shards, names) = four_region_cluster();

        // Write items with text metadata for search.
        for i in 1..=5u64 {
            let eid = EntityId::new(i);
            let mut meta = HashMap::new();
            meta.insert("title".to_string(), format!("jazz piano track {i}"));
            cluster.write_item_with_metadata(eid, &meta).unwrap();
            cluster.write_signal("view", eid, 1.0).unwrap();
        }

        let search_query = tidaldb::query::search::Search::builder()
            .query("jazz")
            .limit(5)
            .build()
            .unwrap();

        let (result, meta) = scatter_gather_search(&cluster, &search_query, &shards, &names, None)
            .expect("scatter-gather search should succeed");

        // Search may or may not find results depending on text index reload timing,
        // but the scatter-gather itself should not error.
        assert!(!meta.degraded, "should not be degraded");
        assert_eq!(meta.shards_queried, 4);
        // Scores descending if items returned.
        for w in result.items.windows(2) {
            assert!(w[0].score >= w[1].score);
        }
    }

    /// `SimulatedCluster` must be `Sync` for the detached scatter-gather
    /// workers (each holds a shared `Arc` and only reads). A regression that
    /// made it non-`Sync` would break the whole fan-out design.
    #[test]
    fn simulated_cluster_is_sync() {
        fn assert_sync<T: Sync + Send + 'static>() {}
        assert_sync::<SimulatedCluster>();
    }

    /// SCATTER-1: a single hung shard must NOT block the whole gather past its
    /// deadline. The slow shard is reported as degraded; the fast shards still
    /// contribute. Exercises [`dispatch_shards`] directly with a closure that
    /// sleeps one shard well past the deadline.
    #[test]
    fn dispatch_shards_slow_shard_does_not_block_deadline() {
        let (cluster, shards, names) = four_region_cluster();
        let slow_shard = shards[2];
        let deadline = Duration::from_millis(60);

        let started = Instant::now();
        let state = dispatch_shards::<u64, _>(
            &cluster,
            &shards,
            &names,
            deadline,
            move |_cluster, shard| {
                if shard == slow_shard {
                    // Far longer than the deadline — simulates a hung shard.
                    std::thread::sleep(Duration::from_secs(2));
                }
                Ok(ShardOutcome {
                    items: vec![u64::from(shard.0)],
                    total_candidates: 1,
                })
            },
        );
        let elapsed = started.elapsed();

        // The coordinator returned without waiting on the hung shard. Allow a
        // generous ceiling for thread-spawn + scheduling jitter, but it must be
        // far below the 2s the slow worker sleeps.
        assert!(
            elapsed < Duration::from_millis(800),
            "gather blocked on slow shard: took {elapsed:?}"
        );

        // The three fast shards contributed; the slow shard is degraded, not
        // silently dropped.
        assert_eq!(state.shards_queried, 3, "fast shards should contribute");
        assert_eq!(
            state.unavailable_shards,
            vec![shard_name(&names, slow_shard)],
            "slow shard must be reported degraded"
        );
        assert_eq!(state.items.len(), 3);
    }

    /// SCATTER-1: with NO slow shard, all shards report and nothing is degraded
    /// even under a tight-but-sufficient budget.
    #[test]
    fn dispatch_shards_all_report_under_budget() {
        let (cluster, shards, names) = four_region_cluster();
        let state = dispatch_shards::<u64, _>(
            &cluster,
            &shards,
            &names,
            Duration::from_millis(500),
            |_cluster, shard| {
                Ok(ShardOutcome {
                    items: vec![u64::from(shard.0)],
                    total_candidates: 2,
                })
            },
        );
        assert_eq!(state.shards_queried, 4);
        assert!(state.unavailable_shards.is_empty());
        // Each of the 4 shards reported total_candidates == 2 (completion order
        // is non-deterministic, so compare the sorted set, not the sequence).
        let mut totals = state.per_shard_totals.clone();
        totals.sort_unstable();
        assert_eq!(totals, vec![2, 2, 2, 2]);
        assert_eq!(state.per_shard_totals.iter().sum::<usize>(), 8);
        assert_eq!(state.items.len(), 4);
    }

    /// SCATTER-1: a shard whose query errors is degraded; the rest still merge.
    #[test]
    fn dispatch_shards_error_shard_is_degraded() {
        let (cluster, shards, names) = four_region_cluster();
        let bad_shard = shards[1];
        let state = dispatch_shards::<u64, _>(
            &cluster,
            &shards,
            &names,
            Duration::from_millis(500),
            move |_cluster, shard| {
                if shard == bad_shard {
                    Err(ServerError::BadRequest("boom".into()))
                } else {
                    Ok(ShardOutcome {
                        items: vec![u64::from(shard.0)],
                        total_candidates: 1,
                    })
                }
            },
        );
        assert_eq!(state.shards_queried, 3);
        assert_eq!(
            state.unavailable_shards,
            vec![shard_name(&names, bad_shard)]
        );
    }

    fn retrieve_result(entity_id: u64, score: f64) -> RetrieveResult {
        RetrieveResult {
            entity_id: EntityId::new(entity_id),
            score,
            rank: 0,
            signals: Vec::new(),
        }
    }

    /// Wrap a [`RetrieveResult`] as if it were returned by `region`, for the
    /// diversity/merge unit tests that exercise [`Sourced`]-keyed paths.
    fn sourced(region: RegionId, entity_id: u64, score: f64) -> Sourced<RetrieveResult> {
        Sourced {
            region,
            item: retrieve_result(entity_id, score),
        }
    }

    /// `entity_shard` must agree with the engine `ShardRouter::hash` (FNV-1a)
    /// for EVERY entity, so server-side write/read routing can never disagree
    /// with the engine's own mapping. A divergent hash silently sends the same
    /// entity to different shards for writes vs reads.
    #[test]
    fn entity_shard_matches_engine_router() {
        let shards = vec![RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
        let router = ShardRouter::hash(shards.len() as u16).unwrap();
        for i in 0..2000u64 {
            let via_server = entity_shard(EntityId::new(i), &shards);
            let via_engine = router.route(EntityId::new(i));
            assert_eq!(
                via_server.0, via_engine.0,
                "entity {i}: server routed to {via_server:?} but engine routed to {via_engine:?}"
            );
        }
    }

    /// Replicated shards return the SAME entity, so the merge must collapse
    /// duplicates and keep the highest-scoring copy — and report that an
    /// overlap was detected.
    #[test]
    fn dedup_collapses_replicated_copies_keeping_best_score() {
        // Entity 1 returned by three replicas with different scores; entity 2
        // by two replicas; entity 3 once. The best-scoring copy's source region
        // must survive (it is the one a later creator lookup reads from).
        let items = vec![
            sourced(RegionId(0), 1, 0.5),
            sourced(RegionId(1), 1, 0.9), // best for entity 1, from region 1
            sourced(RegionId(2), 1, 0.7),
            sourced(RegionId(0), 2, 0.3),
            sourced(RegionId(3), 2, 0.4), // best for entity 2, from region 3
            sourced(RegionId(2), 3, 0.8),
        ];
        let (deduped, overlap) = dedup_by_entity(items);
        assert!(overlap, "duplicates across replicas must set overlap=true");
        assert_eq!(deduped.len(), 3, "one row per distinct entity");

        let mut by_id: HashMap<u64, (f64, RegionId)> = HashMap::new();
        for s in &deduped {
            by_id.insert(s.item.entity_id.as_u64(), (s.item.score, s.region));
        }
        assert!((by_id[&1].0 - 0.9).abs() < f64::EPSILON, "kept best for e1");
        assert_eq!(by_id[&1].1, RegionId(1), "kept the winning copy's region");
        assert!((by_id[&2].0 - 0.4).abs() < f64::EPSILON, "kept best for e2");
        assert_eq!(by_id[&2].1, RegionId(3), "kept the winning copy's region");
        assert!((by_id[&3].0 - 0.8).abs() < f64::EPSILON);
    }

    /// Disjoint (entity-sharded) shards never return the same entity, so dedup
    /// is a no-op and reports no overlap.
    #[test]
    fn dedup_no_overlap_for_disjoint_shards() {
        let items = vec![
            sourced(RegionId(0), 10, 0.5),
            sourced(RegionId(1), 20, 0.6),
            sourced(RegionId(2), 30, 0.7),
        ];
        let (deduped, overlap) = dedup_by_entity(items);
        assert!(!overlap, "disjoint shards must report overlap=false");
        assert_eq!(deduped.len(), 3);
    }

    /// REPLICATED: the candidate universe is ONE replica's worth, not the sum.
    /// Three replicas each reporting 100 candidates is 100 distinct, not 300.
    #[test]
    fn reconcile_replicated_uses_max_not_sum() {
        let total = reconcile_total_candidates(&[100, 100, 100], 10, /*overlap*/ true);
        assert_eq!(total, 100, "replicated shards must not be summed");
    }

    /// ENTITY-SHARDED: disjoint universes genuinely add up.
    #[test]
    fn reconcile_disjoint_sums() {
        let total = reconcile_total_candidates(&[30, 40, 30], 50, /*overlap*/ false);
        assert_eq!(total, 100, "disjoint shards must be summed");
    }

    /// The reported total can never be smaller than the items actually
    /// returned (defends against a stale/under-reported per-shard total).
    #[test]
    fn reconcile_floors_at_deduped_len() {
        let total = reconcile_total_candidates(&[2, 2], 5, /*overlap*/ true);
        assert_eq!(total, 5, "must floor at the deduped item count");
    }

    /// A single NaN score must NOT scramble the ordering of the real-scored
    /// items. `partial_cmp(...).unwrap_or(Equal)` treated NaN as equal to
    /// everything, degrading the sort to an unstable partial order where one
    /// poisoned candidate could reorder the rest; `total_cmp` gives a total
    /// order so the real scores stay strictly descending and the NaN item is
    /// ranked deterministically — never dropped, and never able to interleave
    /// with the real scores.
    #[test]
    fn merge_and_assemble_nan_score_does_not_scramble_order() {
        let (cluster, _shards, _names) = four_region_cluster();
        let leader = cluster.leader_region();
        // No creator_id → diversity is a no-op and every item is retained, so we
        // observe pure sort behavior. Entity 3 carries a NaN score.
        let items = vec![
            sourced(leader, 1, 0.2),
            sourced(leader, 2, 0.9),
            sourced(leader, 3, f64::NAN),
            sourced(leader, 4, 0.5),
        ];
        let merged = merge_and_assemble(&cluster, items, &[4], None, 10);

        // Nothing dropped (no diversity cap, limit exceeds the set).
        assert_eq!(
            merged.items.len(),
            4,
            "NaN must not drop or duplicate items"
        );
        assert!(merged.constraints_satisfied);

        // The three real-scored items stay strictly descending regardless of
        // where the NaN landed.
        let real: Vec<f64> = merged
            .items
            .iter()
            .map(MergeItem::score)
            .filter(|s| !s.is_nan())
            .collect();
        assert_eq!(
            real,
            vec![0.9, 0.5, 0.2],
            "real scores must stay descending"
        );

        // `total_cmp` ranks a (positive) NaN as the greatest value, so in this
        // descending compare it sorts to the FRONT — deterministically, every
        // run — rather than randomly interleaving among the real scores as the
        // old partial_cmp-as-Equal path allowed.
        assert!(
            merged.items.first().is_some_and(|s| s.score().is_nan()),
            "NaN item must sort deterministically, not scramble the rest"
        );

        // Exactly one NaN survives and it is the only non-finite entry.
        let nan_count = merged.items.iter().filter(|s| s.score().is_nan()).count();
        assert_eq!(nan_count, 1, "the single NaN item is retained exactly once");

        // Ranks are 1-based and contiguous over the whole merged page.
        for (i, item) in merged.items.iter().enumerate() {
            assert_eq!(item.rank, i + 1, "rank must be 1-based contiguous");
        }
    }

    /// The shared merge tail honors the page `limit` and assigns 1-based ranks
    /// for SEARCH items too (proving [`merge_and_assemble`] is generic over the
    /// result type, not just `RetrieveResult`).
    #[test]
    fn merge_and_assemble_truncates_and_ranks_search_items() {
        use tidaldb::query::search::SearchResultItem;

        let (cluster, _shards, _names) = four_region_cluster();
        let leader = cluster.leader_region();
        let make = |entity_id: u64, score: f64| Sourced {
            region: leader,
            item: SearchResultItem {
                entity_id: EntityId::new(entity_id),
                score,
                rank: 0,
                bm25_score: None,
                semantic_score: None,
                signals: Vec::new(),
                metadata: None,
            },
        };
        let items = vec![make(1, 0.1), make(2, 0.9), make(3, 0.5), make(4, 0.7)];
        let merged = merge_and_assemble(&cluster, items, &[4], None, 2);

        assert_eq!(merged.items.len(), 2, "limit must truncate to the page");
        // Top-2 by score descending.
        assert!((merged.items[0].score - 0.9).abs() < f64::EPSILON);
        assert!((merged.items[1].score - 0.7).abs() < f64::EPSILON);
        assert_eq!(merged.items[0].rank, 1);
        assert_eq!(merged.items[1].rank, 2);
    }

    /// End-to-end: across a 4-way REPLICATED cluster, `total_candidates` must
    /// NOT be 4x the single-shard count, and the merged items must not contain
    /// duplicate entities.
    #[test]
    fn scatter_gather_retrieve_total_candidates_not_double_counted() {
        let (cluster, shards, names) = four_region_cluster();
        for i in 1..=12u64 {
            let eid = EntityId::new(i);
            cluster
                .write_item_with_metadata(eid, &HashMap::new())
                .unwrap();
            cluster.write_signal("view", eid, i as f64).unwrap();
        }

        // Single-shard baseline candidate count.
        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(12)
            .build()
            .unwrap();
        let single = cluster.retrieve(shards[0], &retrieve).unwrap();
        let single_total = single.total_candidates;
        assert!(single_total > 0, "baseline shard should see candidates");

        let (result, meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(500)).unwrap();
        assert_eq!(meta.shards_queried, 4);

        // Replicated: merged total must equal one replica's universe, NOT 4x.
        assert_eq!(
            result.total_candidates, single_total,
            "replicated shards double-counted: {} != {single_total}",
            result.total_candidates
        );

        // No duplicate entities survived the merge.
        let mut seen = HashSet::new();
        for item in &result.items {
            assert!(
                seen.insert(item.entity_id.as_u64()),
                "duplicate entity {} in merged result",
                item.entity_id.as_u64()
            );
        }
    }

    /// End-to-end: coordinator-level `max_per_creator` must hold across the
    /// MERGED set, not just per shard. Twelve items all from creator 7; with
    /// `max_per_creator = 3` the merged feed may contain at most 3.
    #[test]
    fn scatter_gather_retrieve_enforces_coordinator_diversity() {
        let (cluster, shards, names) = four_region_cluster();
        for i in 1..=12u64 {
            let eid = EntityId::new(i);
            let mut meta = HashMap::new();
            meta.insert("creator_id".to_string(), "7".to_string());
            cluster.write_item_with_metadata(eid, &meta).unwrap();
            cluster.write_signal("view", eid, i as f64).unwrap();
        }

        let retrieve = tidaldb::query::retrieve::Retrieve::builder()
            .profile("trending")
            .limit(12)
            .diversity(tidaldb::ranking::diversity::DiversityConstraints::new().max_per_creator(3))
            .build()
            .unwrap();

        let (result, _meta) =
            scatter_gather_retrieve(&cluster, &retrieve, &shards, &names, Some(500)).unwrap();

        let from_creator_7 = result.items.len();
        assert!(
            from_creator_7 <= 3,
            "coordinator diversity breached: {from_creator_7} items from one creator (cap 3)"
        );
        assert!(
            !result.constraints_satisfied,
            "dropping items for the cap must clear constraints_satisfied"
        );
    }

    /// A zero `max_per_creator` is treated as "no cap" (matches the engine),
    /// never as "drop everything".
    #[test]
    fn enforce_max_per_creator_zero_is_no_cap() {
        let (cluster, _shards, _names) = four_region_cluster();
        let leader = cluster.leader_region();
        let mut meta = HashMap::new();
        meta.insert("creator_id".to_string(), "7".to_string());
        for i in 1..=3u64 {
            cluster
                .write_item_with_metadata(EntityId::new(i), &meta)
                .unwrap();
        }
        let mut items = vec![
            sourced(leader, 1, 0.9),
            sourced(leader, 2, 0.8),
            sourced(leader, 3, 0.7),
        ];
        let dropped = enforce_max_per_creator(&cluster, &mut items, 0);
        assert_eq!(dropped, 0, "zero cap must not drop anything");
        assert_eq!(items.len(), 3);
    }

    /// Items with no resolvable creator are never dropped by the cap (a missing
    /// `creator_id` can only ever UNDER-enforce, never hide a result).
    #[test]
    fn enforce_max_per_creator_keeps_unattributed_items() {
        let (cluster, _shards, _names) = four_region_cluster();
        let leader = cluster.leader_region();
        // Write items WITHOUT creator_id.
        for i in 1..=5u64 {
            cluster
                .write_item_with_metadata(EntityId::new(i), &HashMap::new())
                .unwrap();
        }
        let mut items: Vec<Sourced<RetrieveResult>> = (1..=5u64)
            .map(|i| sourced(leader, i, 1.0 / i as f64))
            .collect();
        let dropped = enforce_max_per_creator(&cluster, &mut items, 1);
        assert_eq!(dropped, 0, "unattributed items must never be capped");
        assert_eq!(items.len(), 5);
    }

    /// Entity-sharded under-enforcement regression: items owned by different
    /// shards must be capped by resolving each item's creator from the shard
    /// that returned it. Writing the same prolific creator's items to disjoint
    /// shards (so the leader does NOT hold the non-leader ones) and merging them
    /// must still respect the cap — the old leader-only lookup resolved the
    /// non-leader items' creator to `None` and let them through uncapped.
    #[test]
    fn enforce_max_per_creator_resolves_from_owning_shard() {
        let (cluster, _shards, _names) = four_region_cluster();
        let regions = [RegionId(0), RegionId(1), RegionId(2), RegionId(3)];

        // Write 8 items from creator 7, each to a DIFFERENT shard's local store
        // only (no replication) — exactly the entity-sharded layout where the
        // leader holds only its own slice.
        let mut meta = HashMap::new();
        meta.insert("creator_id".to_string(), "7".to_string());
        let mut items: Vec<Sourced<RetrieveResult>> = Vec::new();
        for i in 1..=8u64 {
            let region = regions[(i as usize - 1) % regions.len()];
            cluster
                .node(region)
                .db
                .write_item_with_metadata(EntityId::new(i), &meta)
                .unwrap();
            // Source-tag each item with the shard that "returned" it.
            items.push(sourced(region, i, 1.0 / i as f64));
        }

        let dropped = enforce_max_per_creator(&cluster, &mut items, 3);
        assert_eq!(
            items.len(),
            3,
            "creator cap must hold across shards; got {} items",
            items.len()
        );
        assert_eq!(dropped, 5, "8 items from one creator, cap 3 → drop 5");
    }

    /// Counter-test proving the bug the fix removes: resolving every item's
    /// creator from the LEADER only (the old behavior) under-enforces, because
    /// the leader does not hold the non-leader shards' items.
    #[test]
    fn leader_only_creator_lookup_under_enforces_when_sharded() {
        let (cluster, _shards, _names) = four_region_cluster();
        let regions = [RegionId(0), RegionId(1), RegionId(2), RegionId(3)];
        let leader = cluster.leader_region();

        let mut meta = HashMap::new();
        meta.insert("creator_id".to_string(), "7".to_string());
        let mut leader_only: HashMap<u64, usize> = HashMap::new();
        let mut kept = 0usize;
        for i in 1..=8u64 {
            let region = regions[(i as usize - 1) % regions.len()];
            cluster
                .node(region)
                .db
                .write_item_with_metadata(EntityId::new(i), &meta)
                .unwrap();
            // Old behavior: always read metadata from the leader.
            let creator = cluster
                .node(leader)
                .db
                .get_item_metadata(EntityId::new(i))
                .ok()
                .flatten()
                .and_then(|m| m.get("creator_id").and_then(|c| c.parse::<u64>().ok()));
            let pass = creator.is_none_or(|cid| {
                let c = leader_only.entry(cid).or_insert(0);
                if *c >= 3 {
                    false
                } else {
                    *c += 1;
                    true
                }
            });
            if pass {
                kept += 1;
            }
        }
        // The leader holds only its own ~2 items, so the other ~6 resolve to
        // None and slip through uncapped — proving the silent under-enforcement
        // the per-shard lookup fixes.
        assert!(
            kept > 3,
            "leader-only lookup should under-enforce (kept {kept} > cap 3)"
        );
    }

    /// C16: the shard-worker semaphore must cap the AGGREGATE concurrent shard
    /// workers. With 2 permits, no more than 2 workers may hold a permit at
    /// once even when 8 contend, and every worker eventually completes (no
    /// deadlock, no lost permit on release).
    #[test]
    fn shard_worker_semaphore_caps_concurrency() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let sem = Arc::new(ShardWorkerSemaphore::new(2));
        let live = Arc::new(AtomicUsize::new(0));
        let peak = Arc::new(AtomicUsize::new(0));
        let completed = Arc::new(AtomicUsize::new(0));

        let handles: Vec<_> = (0..8)
            .map(|_| {
                let sem = Arc::clone(&sem);
                let live = Arc::clone(&live);
                let peak = Arc::clone(&peak);
                let completed = Arc::clone(&completed);
                std::thread::spawn(move || {
                    // Generous timeout: every worker should get a permit
                    // eventually (the holders release quickly).
                    let permit = sem
                        .acquire_timeout(Duration::from_secs(5))
                        .expect("permit must become available within timeout");
                    let now = live.fetch_add(1, Ordering::SeqCst) + 1;
                    peak.fetch_max(now, Ordering::SeqCst);
                    // Hold the permit briefly so contention is real.
                    std::thread::sleep(Duration::from_millis(5));
                    live.fetch_sub(1, Ordering::SeqCst);
                    completed.fetch_add(1, Ordering::SeqCst);
                    drop(permit);
                })
            })
            .collect();
        for h in handles {
            h.join().expect("worker thread must not panic");
        }

        assert!(
            peak.load(Ordering::SeqCst) <= 2,
            "no more than 2 workers may hold a permit at once, saw {}",
            peak.load(Ordering::SeqCst)
        );
        assert_eq!(
            completed.load(Ordering::SeqCst),
            8,
            "every worker must complete (no deadlock, permits all returned)"
        );
        // All permits returned: a fresh acquire succeeds immediately.
        assert!(
            sem.acquire_timeout(Duration::from_millis(1)).is_some(),
            "all permits should be back after every worker finished"
        );
    }

    /// C16: a worker that cannot get a permit before its deadline retires
    /// instead of parking forever. With zero spare permits and a tiny timeout,
    /// `acquire_timeout` returns `None` so the caller degrades the shard.
    #[test]
    fn shard_worker_semaphore_times_out_when_saturated() {
        let sem = ShardWorkerSemaphore::new(1);
        let _held = sem
            .acquire_timeout(Duration::from_millis(1))
            .expect("first acquire takes the only permit");
        // No permit left; a short-deadline acquire must give up (degrade),
        // not block indefinitely.
        let denied = sem.acquire_timeout(Duration::from_millis(10));
        assert!(
            denied.is_none(),
            "saturated semaphore must time out so the worker retires and degrades"
        );
    }
}
