drasi-lib 0.9.2

// Copyright 2025 The Drasi Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! In-memory query output state with O(1) result-set operations and outbox ring buffer.
//!
//! `QueryOutputState` replaces the naive `Vec<serde_json::Value>` approach with an
//! `im::HashMap` keyed by `row_signature`, providing:
//! - O(1) insert, update, and delete operations
//! - O(1) structural-sharing clones for non-blocking snapshot reads
//! - A bounded ring buffer (`outbox`) of recent `QueryResult` emissions

use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;

use serde::{Deserialize, Serialize};
use tokio_stream::Stream;

use crate::channels::{QueryResult, ResultDiff};

/// Wire envelope pairing a row's engine-stamped `row_signature` with its data.
///
/// Used as the per-row serialization format when snapshot rows cross the FFI
/// boundary, so the canonical row identity survives instead of being dropped.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyedSnapshotRow {
    /// The row's `row_signature` (engine-computed identity). `0` means unknown.
    pub k: u64,
    /// The row's data payload as a JSON object (the query result columns and values).
    pub v: serde_json::Value,
}

/// Default outbox capacity if not configured.
pub const DEFAULT_OUTBOX_CAPACITY: usize = 1000;

/// In-memory state tracking the live result set and recent emissions for a query.
///
/// This struct is held behind `Arc<RwLock<...>>` in `DrasiQuery`. Writers acquire
/// a write lock to apply diffs and push to the outbox. Readers (e.g., `fetch_snapshot`)
/// acquire a read lock and clone the `im::HashMap` in O(1) via structural sharing.
///
/// All fields are private to enforce invariants (sequence monotonicity, ring buffer
/// bounds). Use accessor methods for read access and `apply_diffs` /
/// `advance_sequence_and_push` for mutations. Persistent queries hydrate this
/// struct from durable storage on start (`hydrate`); after that, reads are
/// memory-served.
#[derive(Debug, Clone)]
pub struct QueryOutputState {
    /// Live result set, keyed by `row_signature` for O(1) updates.
    /// Uses `im::HashMap` for O(1) structural-sharing clones (the clone itself is
    /// constant-time; access still requires the enclosing `RwLock` read lock).
    results: im::HashMap<u64, serde_json::Value>,
    /// The result sequence number the snapshot reflects.
    /// Incremented only when non-empty diffs are emitted.
    as_of_sequence: u64,
    /// Ring buffer of recent `QueryResult` emissions (bounded by `outbox_capacity`).
    /// Stored as `Arc` for zero-copy dispatch to reactions.
    outbox: VecDeque<Arc<QueryResult>>,
    /// Maximum number of entries retained in the outbox.
    outbox_capacity: usize,
    /// True after startup hydrate/reset. Distinct from `as_of_sequence > 0`
    /// because bootstrap can fill `results` while the sequence is still 0.
    initialized: bool,
    /// Bumped on each output wipe/rebuild so reactions can tell sequence 1
    /// of a new generation from sequence 1 of a previous one.
    generation: u64,
}

impl QueryOutputState {
    /// Maximum allowed outbox capacity to prevent memory exhaustion from misconfiguration.
    const MAX_OUTBOX_CAPACITY: usize = 1_000_000;

    /// Create a new empty `QueryOutputState` with the given outbox capacity.
    ///
    /// A capacity of 0 is treated as 1 (at least one entry must be retainable
    /// for correct dispatch semantics). Values above `MAX_OUTBOX_CAPACITY` (1M)
    /// are clamped to prevent unbounded memory growth from misconfiguration.
    pub fn new(outbox_capacity: usize) -> Self {
        let effective_capacity = outbox_capacity.clamp(1, Self::MAX_OUTBOX_CAPACITY);
        Self {
            results: im::HashMap::new(),
            as_of_sequence: 0,
            // Pre-allocate up to 1024 slots; the deque grows automatically for larger capacities.
            outbox: VecDeque::with_capacity(effective_capacity.min(1024)),
            outbox_capacity: effective_capacity,
            initialized: false,
            generation: 0,
        }
    }

    /// Apply a set of result diffs to the live result set using O(1) HashMap operations.
    ///
    /// This does NOT increment the sequence or push to the outbox — that is done
    /// separately by the caller after constructing the `QueryResult`.
    pub fn apply_diffs(&mut self, diffs: &[ResultDiff]) {
        for diff in diffs {
            match diff {
                ResultDiff::Add {
                    data,
                    row_signature,
                } => {
                    self.results.insert(*row_signature, data.clone());
                }
                ResultDiff::Delete { row_signature, .. } => {
                    self.results.remove(row_signature);
                }
                ResultDiff::Update {
                    after,
                    row_signature,
                    ..
                } => {
                    self.results.insert(*row_signature, after.clone());
                }
                ResultDiff::Aggregation {
                    after,
                    row_signature,
                    ..
                } => {
                    // Insert/overwrite the aggregation result for this group.
                    // Note: identity-value detection (empty group removal) depends on #384
                    // and will be handled in a follow-up.
                    self.results.insert(*row_signature, after.clone());
                }
                ResultDiff::Noop => {}
            }
        }
    }

    /// Increment the sequence counter, wrap the result in an `Arc`, push to the outbox,
    /// and return the `Arc<QueryResult>` for zero-copy dispatch.
    /// Evicts the oldest entry if the outbox is at capacity.
    pub fn advance_sequence_and_push(&mut self, mut result: QueryResult) -> Arc<QueryResult> {
        self.as_of_sequence = self.as_of_sequence.saturating_add(1);
        result.sequence = self.as_of_sequence;

        let arc_result = Arc::new(result);
        self.push_outbox(arc_result.clone());
        arc_result
    }

    /// Apply diffs and install a sequence that was already staged to durable storage.
    ///
    /// Unlike [`advance_sequence_and_push`], this does not independently increment
    /// `as_of_sequence`. The caller must pass the sequence written inside the
    /// session transaction so in-memory state cannot diverge from durable output.
    /// In-memory sequence still only moves after that durable commit succeeds.
    pub fn apply_committed_sequence(
        &mut self,
        sequence: u64,
        diffs: &[ResultDiff],
        mut result: QueryResult,
    ) -> Arc<QueryResult> {
        let expected = self.as_of_sequence.saturating_add(1);
        if sequence != expected {
            log::error!(
                "committed output sequence {sequence} != next in-memory sequence {expected}"
            );
        }
        self.apply_diffs(diffs);
        self.as_of_sequence = sequence;
        result.sequence = sequence;
        let arc_result = Arc::new(result);
        self.push_outbox(arc_result.clone());
        arc_result
    }

    fn push_outbox(&mut self, arc_result: Arc<QueryResult>) {
        if self.outbox.len() >= self.outbox_capacity {
            self.outbox.pop_front();
        }
        self.outbox.push_back(arc_result);
    }

    /// Return the live result set as a `Vec` for backward compatibility with `get_current_results`.
    pub fn get_results_as_vec(&self) -> Vec<serde_json::Value> {
        self.results.values().cloned().collect()
    }

    /// Return the current outbox capacity.
    pub fn outbox_capacity(&self) -> usize {
        self.outbox_capacity
    }

    /// Return the current sequence number.
    pub fn as_of_sequence(&self) -> u64 {
        self.as_of_sequence
    }

    /// Return the number of entries currently in the outbox.
    pub fn outbox_len(&self) -> usize {
        self.outbox.len()
    }

    /// Return the sequence of the earliest entry in the outbox, or `None` if empty.
    pub fn outbox_earliest_seq(&self) -> Option<u64> {
        self.outbox.front().map(|r| r.sequence)
    }

    /// Return the number of results in the live result set.
    pub fn results_len(&self) -> usize {
        self.results.len()
    }

    /// Get a result by its row signature.
    pub fn get_result(&self, row_signature: &u64) -> Option<&serde_json::Value> {
        self.results.get(row_signature)
    }

    /// Clone the live result set as an `im::HashMap` (O(1) via structural sharing).
    ///
    /// This is used by `DrasiQuery::fetch_snapshot()` to take a lightweight clone
    /// under the read lock, then build the `SnapshotResponse` outside the lock.
    pub fn clone_results(&self) -> im::HashMap<u64, serde_json::Value> {
        self.results.clone()
    }

    /// Fetch outbox entries after the given sequence number.
    ///
    /// Returns `Ok(entries)` if the requested position is available in the ring buffer,
    /// or `Err(OutboxGap)` if the position has been evicted.
    pub fn fetch_outbox_after(
        &self,
        after_sequence: u64,
    ) -> Result<Vec<Arc<QueryResult>>, OutboxGap> {
        if after_sequence >= self.as_of_sequence {
            // Caller is at or ahead of the current sequence — nothing to return
            return Ok(Vec::new());
        }

        // Find the earliest sequence in the outbox
        let earliest = self
            .outbox
            .front()
            .map(|r| r.sequence)
            .unwrap_or(self.as_of_sequence + 1);

        if after_sequence + 1 < earliest {
            return Err(OutboxGap {
                requested: after_sequence,
                earliest_available: earliest,
                latest_sequence: self.as_of_sequence,
                config_hash: 0, // Enriched by `DrasiQuery::fetch_outbox()`
            });
        }

        // Collect entries with sequence > after_sequence (cheap Arc clone)
        let entries: Vec<Arc<QueryResult>> = self
            .outbox
            .iter()
            .filter(|r| r.sequence > after_sequence)
            .cloned()
            .collect();

        Ok(entries)
    }

    /// Populate this state from durable live rows, outbox entries, and a reconciled
    /// high-water sequence.
    ///
    /// The sequence is never lowered: `as_of_sequence` becomes
    /// `max(current, as_of_sequence)`. Outbox entries are sorted and trimmed to
    /// `outbox_capacity`, keeping the newest. Callers must validate durable
    /// consistency with [`reconcile_durable_output`] before invoking this.
    ///
    /// Hydrate is a startup-only operation from uninitialized state. Calling it
    /// after [`initialized`](Self::initialized) is a programming error; debug
    /// builds assert.
    pub fn hydrate(
        &mut self,
        results: im::HashMap<u64, serde_json::Value>,
        mut outbox: Vec<Arc<QueryResult>>,
        as_of_sequence: u64,
        generation: u64,
    ) {
        debug_assert!(
            !self.initialized,
            "hydrate must run once from uninitialized QueryOutputState"
        );
        outbox.sort_by_key(|result| result.sequence);
        if outbox.len() > self.outbox_capacity {
            let skip = outbox.len() - self.outbox_capacity;
            outbox = outbox.split_off(skip);
        }

        self.results = results;
        self.outbox = VecDeque::from(outbox);
        self.as_of_sequence = self.as_of_sequence.max(as_of_sequence);
        self.generation = generation;
        self.initialized = true;
    }

    /// Clear live rows, outbox, and sequence (used by AutoReset output wipe).
    ///
    /// Bumps [`generation`](Self::generation) so reactions do not treat the new
    /// sequence 1 as a duplicate of the previous generation.
    pub fn reset(&mut self) {
        self.reset_from_generation(self.generation);
    }

    /// Like [`reset`](Self::reset), but bump from `persisted` instead of RAM.
    ///
    /// Use this when hydrate never ran, so in-memory generation is still 0
    /// while disk may already hold a higher value.
    pub fn reset_from_generation(&mut self, persisted: u64) {
        self.results.clear();
        self.outbox.clear();
        self.as_of_sequence = 0;
        self.generation = next_output_generation(persisted, self.generation);
        self.initialized = true;
    }

    /// Mark output as initialized without changing rows or sequence.
    ///
    /// Used after a first-run start that skipped durable hydrate so a later
    /// same-process stop/start does not overwrite bootstrap results.
    pub fn mark_initialized(&mut self) {
        self.initialized = true;
    }

    /// Whether startup hydrate (or a wipe/rebuild) has already initialized this state.
    pub fn initialized(&self) -> bool {
        self.initialized
    }

    /// Output generation. `0` until the first wipe/rebuild.
    pub fn generation(&self) -> u64 {
        self.generation
    }
}

/// Mix query identity hash with output generation.
///
/// Generation 0 returns `config_hash` unchanged so existing reaction
/// checkpoints keep matching. A wipe/rebuild changes the value reactions
/// compare, so they apply recovery policy instead of skipping 1..=N.
pub fn output_epoch_hash(config_hash: u64, generation: u64) -> u64 {
    config_hash.wrapping_add(generation.wrapping_mul(0x9E3779B97F4A7C15))
}

/// Next output generation after an AutoReset, including resume of a crash
/// mid-wipe.
///
/// Always bump from `max(disk, RAM)`. Do not subtract 1 from a persisted
/// value: a crash after the in-progress marker and before the generation
/// write still has the old generation on disk, and rewriting it lets a
/// reaction skip the new sequence 1.
pub(crate) fn next_output_generation(persisted: u64, ram: u64) -> u64 {
    persisted.max(ram).saturating_add(1)
}

/// Inconsistency between durable result sequence, outbox, and live rows.
///
/// On query start this is a `Strict` failure or an `AutoReset` output wipe.
/// Never continue processing with sequence 0 against non-empty durable output.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum DurableOutputInconsistency {
    /// Durable outbox high-water is ahead of the stored result sequence.
    #[error("outbox high-water {outbox_hwm} is ahead of stored result sequence {stored_sequence}")]
    OutboxAheadOfSequence {
        stored_sequence: u64,
        outbox_hwm: u64,
    },
    /// Retained outbox keys are not a contiguous sequence window.
    #[error("durable outbox has a gap; retained sequences: {retained:?}")]
    GappedOutbox { retained: Vec<u64> },
    /// Stored sequence is > 0 but live rows could not be read.
    #[error("stored result sequence {stored_sequence} has no readable live rows")]
    MissingLiveRows { stored_sequence: u64 },
    /// Sequence is 0/absent but durable live rows exist.
    #[error(
        "sequence 0 against non-empty durable output (live_rows={live_rows}, outbox_hwm={outbox_hwm:?})"
    )]
    SequenceZeroAgainstDurableOutput {
        live_rows: usize,
        outbox_hwm: Option<u64>,
    },
    /// A durable outbox payload could not be deserialized.
    #[error("failed to deserialize durable outbox entry at sequence {sequence}: {message}")]
    CorruptOutbox { sequence: u64, message: String },
    /// A durable live-results row could not be deserialized.
    #[error("failed to deserialize durable live row {row_signature}: {message}")]
    CorruptLiveRow { row_signature: u64, message: String },
    /// A durable store could not be read at startup.
    ///
    /// This is a transient I/O failure, not structural corruption. Callers must
    /// fail start (retryable) and must **not** AutoReset/wipe.
    #[error("failed to read durable query output: {message}")]
    ReadFailed { message: String },
}

impl DurableOutputInconsistency {
    /// Transient storage read failures must not authorize a destructive wipe.
    pub(crate) fn is_transient_read(&self) -> bool {
        matches!(self, Self::ReadFailed { .. })
    }
}

/// Reconcile durable output high-water marks without lowering either side.
///
/// Returns the sequence to install into `QueryOutputState`. The next emitted
/// result must be this value plus one.
///
/// # Errors
/// Returns [`DurableOutputInconsistency`] if the stored sequence, outbox
/// high-water mark, and live row count are not mutually consistent (e.g. a
/// gapped or duplicate outbox, an outbox ahead of the stored sequence, missing
/// live rows for a stored sequence, or non-empty durable output with sequence 0).
pub(crate) fn reconcile_durable_output(
    stored_sequence: Option<u64>,
    outbox_sequences: &[u64],
    live_row_count: usize,
    live_rows_readable: bool,
) -> Result<u64, DurableOutputInconsistency> {
    let stored = stored_sequence.unwrap_or(0);

    if !live_rows_readable && stored > 0 {
        return Err(DurableOutputInconsistency::MissingLiveRows {
            stored_sequence: stored,
        });
    }

    let mut retained: Vec<u64> = outbox_sequences.to_vec();
    retained.sort_unstable();
    if !retained.is_empty() {
        let unique_count = {
            let mut deduped = retained.clone();
            deduped.dedup();
            deduped.len()
        };
        if unique_count != retained.len()
            || retained.windows(2).any(|window| window[1] != window[0] + 1)
        {
            return Err(DurableOutputInconsistency::GappedOutbox { retained });
        }
    }

    let outbox_hwm = retained.last().copied();
    if stored == 0 && (live_row_count > 0 || outbox_hwm.is_some()) {
        return Err(
            DurableOutputInconsistency::SequenceZeroAgainstDurableOutput {
                live_rows: live_row_count,
                outbox_hwm,
            },
        );
    }

    if let Some(hwm) = outbox_hwm {
        if hwm > stored {
            return Err(DurableOutputInconsistency::OutboxAheadOfSequence {
                stored_sequence: stored,
                outbox_hwm: hwm,
            });
        }
    }

    // stored >= outbox_hwm (or outbox empty). Do not lower to the outbox HWM.
    Ok(stored)
}

/// Error returned when the requested outbox position has been evicted.
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
#[error("Outbox gap: requested after seq {requested}, but earliest available is {earliest_available} (latest: {latest_sequence})")]
pub struct OutboxGap {
    /// The sequence the caller requested (wants entries after this).
    pub requested: u64,
    /// The earliest sequence still available in the outbox.
    pub earliest_available: u64,
    /// The latest sequence in the outbox.
    pub latest_sequence: u64,
    /// The query's config hash (set by `DrasiQuery`, not by `QueryOutputState`).
    pub config_hash: u64,
}

/// Error returned by `fetch_snapshot` or `fetch_outbox` when the query is not
/// in a state that can serve the request.
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
pub enum FetchError {
    /// The query finished bootstrapping but ended in a non-Running state
    /// (e.g., Error, Stopped). The snapshot/outbox may be incomplete.
    #[error("Query is not running (status: {status:?})")]
    NotRunning {
        status: crate::channels::ComponentStatus,
    },
    /// The bootstrap did not complete within the allowed timeout.
    #[error("Timed out waiting for query to finish bootstrapping")]
    TimedOut,
    /// (fetch_outbox only) The requested outbox position has been evicted.
    #[error(transparent)]
    OutboxGap(#[from] OutboxGap),
}

/// Response from `fetch_snapshot` on the Query trait.
///
/// Contains the live result set as an O(1) `im::HashMap` clone (private) and
/// exposes it via `stream()`, `to_vec()`, `len()`, and `is_empty()`.
#[derive(Debug, Clone)]
pub struct SnapshotResponse {
    /// The live result set (private — not part of the public API).
    results: im::HashMap<u64, serde_json::Value>,
    /// The sequence number this snapshot reflects.
    pub as_of_sequence: u64,
    /// The query's configuration hash at the time of the snapshot.
    pub config_hash: u64,
    /// Output generation at the time of the snapshot. `0` until a wipe/rebuild.
    pub output_generation: u64,
}

impl SnapshotResponse {
    /// Create a new `SnapshotResponse` from an `im::HashMap` clone.
    pub fn new(
        results: im::HashMap<u64, serde_json::Value>,
        as_of_sequence: u64,
        config_hash: u64,
    ) -> Self {
        Self {
            results,
            as_of_sequence,
            config_hash,
            output_generation: 0,
        }
    }

    /// Set the output generation carried with this snapshot.
    pub fn with_output_generation(mut self, generation: u64) -> Self {
        self.output_generation = generation;
        self
    }

    /// Return an async stream of the result values.
    ///
    /// The stream yields each `serde_json::Value` from the underlying `im::HashMap`
    /// without holding any lock (the clone was taken under the read lock).
    pub fn stream(self) -> impl Stream<Item = serde_json::Value> + Send {
        tokio_stream::iter(self.results.into_iter().map(|(_, v)| v))
    }

    /// Return an async stream of `(row_signature, value)` pairs.
    ///
    /// Like [`stream`](Self::stream) but preserves each row's `row_signature`
    /// (the canonical identity) instead of dropping it.
    pub fn stream_keyed(self) -> impl Stream<Item = (u64, serde_json::Value)> + Send {
        tokio_stream::iter(self.results)
    }

    /// Collect the results into a `Vec<serde_json::Value>`.
    pub fn to_vec(&self) -> Vec<serde_json::Value> {
        self.results.values().cloned().collect()
    }

    /// Return the number of results in the snapshot.
    pub fn len(&self) -> usize {
        self.results.len()
    }

    /// Return `true` if the snapshot contains no results.
    pub fn is_empty(&self) -> bool {
        self.results.is_empty()
    }
}

/// Response from `fetch_outbox` on the Query trait.
#[derive(Debug, Clone)]
pub struct OutboxResponse {
    /// The contiguous set of `QueryResult` entries after the requested position.
    pub results: Vec<Arc<QueryResult>>,
    /// The latest sequence number in the query's output state.
    pub latest_sequence: u64,
    /// The query's configuration hash.
    pub config_hash: u64,
    /// Output generation. `0` until a wipe/rebuild.
    pub output_generation: u64,
}

/// Streaming snapshot response for the bootstrap path.
///
/// Wraps a `Stream` of `serde_json::Value` rows plus snapshot metadata.
/// Created from either an in-process `SnapshotResponse` (via `from_snapshot()`)
/// or an FFI iterator (via the plugin SDK).
///
/// Consumers iterate with `while let Some(row) = stream.next().await { ... }`
/// or call `collect_vec()` to drain all rows into a `Vec`.
pub struct SnapshotStream {
    inner: Pin<Box<dyn Stream<Item = (u64, serde_json::Value)> + Send>>,
    /// The sequence number this snapshot reflects.
    pub as_of_sequence: u64,
    /// The query's configuration hash at the time of the snapshot.
    pub config_hash: u64,
}

impl SnapshotStream {
    /// Create a `SnapshotStream` from an in-process `SnapshotResponse`.
    pub fn from_snapshot(snapshot: SnapshotResponse) -> Self {
        let as_of_sequence = snapshot.as_of_sequence;
        let config_hash = snapshot.config_hash;
        Self {
            inner: Box::pin(snapshot.stream_keyed()),
            as_of_sequence,
            config_hash,
        }
    }

    /// Create a `SnapshotStream` from a stream of bare row values.
    ///
    /// Rows created this way have an unknown `row_signature` (stamped as `0`);
    /// use [`from_keyed_stream`](Self::from_keyed_stream) when signatures are
    /// available.
    pub fn from_stream(
        stream: impl Stream<Item = serde_json::Value> + Send + 'static,
        as_of_sequence: u64,
        config_hash: u64,
    ) -> Self {
        use tokio_stream::StreamExt;
        Self {
            inner: Box::pin(stream.map(|v| (0u64, v))),
            as_of_sequence,
            config_hash,
        }
    }

    /// Create a `SnapshotStream` from a stream of `(row_signature, value)` pairs.
    ///
    /// Prefer this over [`from_stream`](Self::from_stream) when row signatures are
    /// available; signatures let downstream consumers deduplicate and match rows
    /// by canonical identity (see `row_signature`).
    pub fn from_keyed_stream(
        stream: impl Stream<Item = (u64, serde_json::Value)> + Send + 'static,
        as_of_sequence: u64,
        config_hash: u64,
    ) -> Self {
        Self {
            inner: Box::pin(stream),
            as_of_sequence,
            config_hash,
        }
    }

    /// Collect all rows from the stream into a `Vec`, dropping signatures.
    ///
    /// Use [`collect_keyed_vec`](Self::collect_keyed_vec) to retain each row's
    /// `row_signature` alongside its data.
    pub async fn collect_vec(self) -> Vec<serde_json::Value> {
        use tokio_stream::StreamExt;
        self.inner.map(|(_, v)| v).collect().await
    }

    /// Collect all rows from the stream into a `Vec` of `(row_signature, value)` pairs.
    pub async fn collect_keyed_vec(self) -> Vec<(u64, serde_json::Value)> {
        use tokio_stream::StreamExt;
        self.inner.collect().await
    }

    /// Collect up to `limit` `(row_signature, value)` pairs, then stop pulling.
    ///
    /// Bounds peak memory when draining a potentially large snapshot into a
    /// capped consumer (e.g. a store with a per-query row limit).
    pub async fn collect_keyed_vec_capped(self, limit: usize) -> Vec<(u64, serde_json::Value)> {
        use tokio_stream::StreamExt;
        self.inner.take(limit).collect().await
    }

    /// Pull the next `(row_signature, value)` pair from the stream.
    pub async fn next_keyed(&mut self) -> Option<(u64, serde_json::Value)> {
        use tokio_stream::StreamExt;
        self.inner.next().await
    }
}

/// Yields `serde_json::Value` rows, **dropping each row's `row_signature`**.
///
/// This trait impl is retained for value-only consumers (e.g. bootstrap replay in
/// other reactions). Identity-preserving consumers must instead use the keyed
/// methods ([`next_keyed`](SnapshotStream::next_keyed) /
/// [`collect_keyed_vec`](SnapshotStream::collect_keyed_vec)); routing snapshot rows
/// through this lossy path is what previously caused the #605 deduplication bug.
impl Stream for SnapshotStream {
    type Item = serde_json::Value;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.inner
            .as_mut()
            .poll_next(cx)
            .map(|opt| opt.map(|(_, v)| v))
    }
}

/// Streaming outbox response for the bootstrap path.
///
/// Wraps a `Stream` of `Arc<QueryResult>` entries plus outbox metadata.
pub struct OutboxStream {
    inner: Pin<Box<dyn Stream<Item = Arc<QueryResult>> + Send>>,
    /// The latest sequence number in the query's output state.
    pub latest_sequence: u64,
    /// The query's configuration hash.
    pub config_hash: u64,
}

impl OutboxStream {
    /// Create an `OutboxStream` from an in-process `OutboxResponse`.
    pub fn from_outbox(outbox: OutboxResponse) -> Self {
        let latest_sequence = outbox.latest_sequence;
        let config_hash = outbox.config_hash;
        Self {
            inner: Box::pin(tokio_stream::iter(outbox.results)),
            latest_sequence,
            config_hash,
        }
    }

    /// Create an `OutboxStream` from an arbitrary `Stream` implementation.
    pub fn from_stream(
        stream: impl Stream<Item = Arc<QueryResult>> + Send + 'static,
        latest_sequence: u64,
        config_hash: u64,
    ) -> Self {
        Self {
            inner: Box::pin(stream),
            latest_sequence,
            config_hash,
        }
    }

    /// Collect all entries from the stream into a `Vec`.
    pub async fn collect_vec(self) -> Vec<Arc<QueryResult>> {
        use tokio_stream::StreamExt;
        self.inner.collect().await
    }
}

impl Stream for OutboxStream {
    type Item = Arc<QueryResult>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.inner.as_mut().poll_next(cx)
    }
}

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

    fn make_query_result(query_id: &str, diffs: Vec<ResultDiff>) -> QueryResult {
        QueryResult::new(
            query_id.to_string(),
            0, // sequence will be set by advance_sequence_and_push
            chrono::Utc::now(),
            diffs,
            HashMap::new(),
        )
    }

    #[test]
    fn test_apply_diffs_add() {
        let mut state = QueryOutputState::new(10);
        let diffs = vec![ResultDiff::Add {
            data: serde_json::json!({"name": "Alice"}),
            row_signature: 100,
        }];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 1);
        assert_eq!(
            state.results.get(&100),
            Some(&serde_json::json!({"name": "Alice"}))
        );
    }

    #[test]
    fn test_apply_diffs_delete() {
        let mut state = QueryOutputState::new(10);
        state
            .results
            .insert(100, serde_json::json!({"name": "Alice"}));

        let diffs = vec![ResultDiff::Delete {
            data: serde_json::json!({"name": "Alice"}),
            row_signature: 100,
        }];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 0);
    }

    #[test]
    fn test_apply_diffs_update() {
        let mut state = QueryOutputState::new(10);
        state
            .results
            .insert(100, serde_json::json!({"name": "Alice"}));

        let diffs = vec![ResultDiff::Update {
            data: serde_json::json!({"name": "Bob"}),
            before: serde_json::json!({"name": "Alice"}),
            after: serde_json::json!({"name": "Bob"}),
            grouping_keys: None,
            row_signature: 100,
        }];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 1);
        assert_eq!(
            state.results.get(&100),
            Some(&serde_json::json!({"name": "Bob"}))
        );
    }

    #[test]
    fn test_apply_diffs_aggregation() {
        let mut state = QueryOutputState::new(10);

        let diffs = vec![ResultDiff::Aggregation {
            before: None,
            after: serde_json::json!({"count": 5}),
            row_signature: 200,
        }];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 1);
        assert_eq!(
            state.results.get(&200),
            Some(&serde_json::json!({"count": 5}))
        );

        // Update aggregation
        let diffs = vec![ResultDiff::Aggregation {
            before: Some(serde_json::json!({"count": 5})),
            after: serde_json::json!({"count": 10}),
            row_signature: 200,
        }];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 1);
        assert_eq!(
            state.results.get(&200),
            Some(&serde_json::json!({"count": 10}))
        );
    }

    #[test]
    fn test_apply_diffs_noop() {
        let mut state = QueryOutputState::new(10);
        state
            .results
            .insert(100, serde_json::json!({"name": "Alice"}));

        let diffs = vec![ResultDiff::Noop];
        state.apply_diffs(&diffs);

        assert_eq!(state.results.len(), 1);
    }

    #[test]
    fn test_advance_sequence_and_push() {
        let mut state = QueryOutputState::new(3);

        let result = make_query_result("q1", vec![]);
        let arc = state.advance_sequence_and_push(result);
        assert_eq!(arc.sequence, 1);
        assert_eq!(state.as_of_sequence, 1);
        assert_eq!(state.outbox.len(), 1);
        assert_eq!(state.outbox.back().unwrap().sequence, 1);

        let result = make_query_result("q1", vec![]);
        let arc = state.advance_sequence_and_push(result);
        assert_eq!(arc.sequence, 2);
        assert_eq!(state.outbox.len(), 2);
    }

    #[test]
    fn test_apply_committed_sequence_uses_durable_seq() {
        let mut state = QueryOutputState::new(3);
        let diffs = vec![ResultDiff::Add {
            data: serde_json::json!({"name": "Alice"}),
            row_signature: 1,
        }];
        let result = make_query_result("q1", diffs.clone());
        let arc = state.apply_committed_sequence(1, &diffs, result);
        assert_eq!(arc.sequence, 1);
        assert_eq!(state.as_of_sequence(), 1);
        assert_eq!(state.results_len(), 1);
        assert_eq!(state.outbox_len(), 1);
    }

    #[test]
    fn test_outbox_capacity_eviction() {
        let mut state = QueryOutputState::new(3);

        for _ in 0..5 {
            let result = make_query_result("q1", vec![]);
            state.advance_sequence_and_push(result);
        }

        assert_eq!(state.outbox.len(), 3);
        assert_eq!(state.as_of_sequence, 5);
        // Oldest should be seq 3 (1 and 2 evicted)
        assert_eq!(state.outbox.front().unwrap().sequence, 3);
        assert_eq!(state.outbox.back().unwrap().sequence, 5);
    }

    #[test]
    fn test_fetch_outbox_after_caught_up() {
        let mut state = QueryOutputState::new(10);

        let result = make_query_result("q1", vec![]);
        state.advance_sequence_and_push(result);

        // Asking for entries after current sequence → empty
        let entries = state.fetch_outbox_after(1).unwrap();
        assert!(entries.is_empty());

        // Asking for entries after a future sequence → also empty
        let entries = state.fetch_outbox_after(100).unwrap();
        assert!(entries.is_empty());
    }

    #[test]
    fn test_fetch_outbox_after_returns_entries() {
        let mut state = QueryOutputState::new(10);

        for _ in 0..5 {
            let result = make_query_result("q1", vec![]);
            state.advance_sequence_and_push(result);
        }

        let entries = state.fetch_outbox_after(2).unwrap();
        assert_eq!(entries.len(), 3);
        assert_eq!(entries[0].sequence, 3);
        assert_eq!(entries[1].sequence, 4);
        assert_eq!(entries[2].sequence, 5);
    }

    #[test]
    fn test_fetch_outbox_after_gap_error() {
        let mut state = QueryOutputState::new(3);

        for _ in 0..5 {
            let result = make_query_result("q1", vec![]);
            state.advance_sequence_and_push(result);
        }

        // Outbox contains seq 3, 4, 5. Requesting after seq 0 → gap
        let err = state.fetch_outbox_after(0).unwrap_err();
        assert_eq!(err.requested, 0);
        assert_eq!(err.earliest_available, 3);
        assert_eq!(err.latest_sequence, 5);
        assert_eq!(err.config_hash, 0); // Default; enriched by DrasiQuery
    }

    #[test]
    fn test_get_results_as_vec() {
        let mut state = QueryOutputState::new(10);
        state.results.insert(1, serde_json::json!({"a": 1}));
        state.results.insert(2, serde_json::json!({"b": 2}));

        let vec = state.get_results_as_vec();
        assert_eq!(vec.len(), 2);
        // Order is not guaranteed from HashMap, just check both values are present
        assert!(vec.contains(&serde_json::json!({"a": 1})));
        assert!(vec.contains(&serde_json::json!({"b": 2})));
    }

    #[test]
    fn test_snapshot_clone_is_independent() {
        let mut state = QueryOutputState::new(10);
        state
            .results
            .insert(1, serde_json::json!({"name": "Alice"}));

        // Clone the results (simulating a snapshot read)
        let snapshot = state.results.clone();

        // Mutate the original
        state.results.insert(1, serde_json::json!({"name": "Bob"}));

        // Snapshot is unchanged (structural sharing)
        assert_eq!(
            snapshot.get(&1),
            Some(&serde_json::json!({"name": "Alice"}))
        );
        assert_eq!(
            state.results.get(&1),
            Some(&serde_json::json!({"name": "Bob"}))
        );
    }

    #[test]
    fn test_outbox_capacity_zero_clamped_to_one() {
        let mut state = QueryOutputState::new(0);
        // Capacity 0 is clamped to 1
        assert_eq!(state.outbox_capacity, 1);

        let result = make_query_result("q1", vec![]);
        state.advance_sequence_and_push(result);
        assert_eq!(state.outbox.len(), 1);

        // Second push evicts the first (capacity is 1)
        let result = make_query_result("q1", vec![]);
        state.advance_sequence_and_push(result);
        assert_eq!(state.outbox.len(), 1);
        assert_eq!(state.outbox.front().unwrap().sequence, 2);
    }

    #[tokio::test]
    async fn snapshot_stream_yields_all_values() {
        use tokio_stream::StreamExt;

        let mut map = im::HashMap::new();
        map.insert(1, serde_json::json!({"id": 1}));
        map.insert(2, serde_json::json!({"id": 2}));
        map.insert(3, serde_json::json!({"id": 3}));

        let snap = SnapshotResponse::new(map, 10, 42);
        assert_eq!(snap.len(), 3);
        assert!(!snap.is_empty());

        // Consume via stream()
        let mut collected: Vec<serde_json::Value> = snap.stream().collect().await;
        collected.sort_by_key(|v| v["id"].as_u64().unwrap());
        assert_eq!(collected.len(), 3);
        assert_eq!(collected[0]["id"], 1);
        assert_eq!(collected[1]["id"], 2);
        assert_eq!(collected[2]["id"], 3);
    }

    #[tokio::test]
    async fn test_snapshot_stream_preserves_row_signatures() {
        use tokio_stream::StreamExt;

        let mut map = im::HashMap::new();
        map.insert(11u64, serde_json::json!({"id": 1}));
        map.insert(22u64, serde_json::json!({"id": 2}));

        let snap = SnapshotResponse::new(map, 5, 7);

        // stream_keyed yields (row_signature, value) pairs.
        let mut keyed: Vec<(u64, serde_json::Value)> = snap.clone().stream_keyed().collect().await;
        keyed.sort_by_key(|(sig, _)| *sig);
        assert_eq!(keyed.len(), 2);
        assert_eq!(keyed[0].0, 11);
        assert_eq!(keyed[0].1["id"], 1);
        assert_eq!(keyed[1].0, 22);

        // SnapshotStream::collect_keyed_vec preserves signatures too.
        let stream = SnapshotStream::from_snapshot(snap);
        assert_eq!(stream.as_of_sequence, 5);
        let mut via_stream = stream.collect_keyed_vec().await;
        via_stream.sort_by_key(|(sig, _)| *sig);
        assert_eq!(via_stream.len(), 2);
        assert_eq!(via_stream[0].0, 11);
        assert_eq!(via_stream[1].0, 22);
    }

    #[tokio::test]
    async fn test_snapshot_stream_from_bare_values_uses_zero_signature() {
        let stream = SnapshotStream::from_stream(
            tokio_stream::iter(vec![serde_json::json!({"id": 1})]),
            0,
            0,
        );
        let keyed = stream.collect_keyed_vec().await;
        assert_eq!(keyed.len(), 1);
        assert_eq!(
            keyed[0].0, 0,
            "bare-value stream rows have unknown signature 0"
        );
    }

    #[test]
    fn reconcile_matching_sequence_and_outbox_hwm() {
        let seq = reconcile_durable_output(Some(4), &[1, 2, 3, 4], 3, true).unwrap();
        assert_eq!(seq, 4);
    }

    #[test]
    fn reconcile_does_not_lower_stored_sequence_below_outbox_hwm() {
        let seq = reconcile_durable_output(Some(5), &[3, 4], 1, true).unwrap();
        assert_eq!(seq, 5);
    }

    #[test]
    fn reconcile_outbox_ahead_of_sequence_is_inconsistent() {
        let err = reconcile_durable_output(Some(4), &[1, 2, 3, 4, 5], 3, true).unwrap_err();
        assert_eq!(
            err,
            DurableOutputInconsistency::OutboxAheadOfSequence {
                stored_sequence: 4,
                outbox_hwm: 5,
            }
        );
    }

    #[test]
    fn reconcile_gapped_outbox_is_inconsistent() {
        let err = reconcile_durable_output(Some(5), &[1, 2, 4, 5], 2, true).unwrap_err();
        assert!(matches!(
            err,
            DurableOutputInconsistency::GappedOutbox { retained } if retained == vec![1, 2, 4, 5]
        ));
    }

    #[test]
    fn reconcile_duplicate_outbox_sequence_is_inconsistent() {
        let err = reconcile_durable_output(Some(3), &[1, 2, 2, 3], 2, true).unwrap_err();
        assert!(matches!(
            err,
            DurableOutputInconsistency::GappedOutbox { retained } if retained == vec![1, 2, 2, 3]
        ));
    }

    #[test]
    fn reconcile_missing_live_rows_for_stored_sequence() {
        let err = reconcile_durable_output(Some(3), &[1, 2, 3], 0, false).unwrap_err();
        assert_eq!(
            err,
            DurableOutputInconsistency::MissingLiveRows { stored_sequence: 3 }
        );
    }

    #[test]
    fn reconcile_sequence_zero_against_live_rows_is_inconsistent() {
        let err = reconcile_durable_output(None, &[], 2, true).unwrap_err();
        assert_eq!(
            err,
            DurableOutputInconsistency::SequenceZeroAgainstDurableOutput {
                live_rows: 2,
                outbox_hwm: None,
            }
        );
    }

    #[test]
    fn reconcile_sequence_zero_against_outbox_is_inconsistent() {
        let err = reconcile_durable_output(None, &[1, 2, 3], 0, true).unwrap_err();
        assert_eq!(
            err,
            DurableOutputInconsistency::SequenceZeroAgainstDurableOutput {
                live_rows: 0,
                outbox_hwm: Some(3),
            }
        );
    }

    #[test]
    fn reconcile_empty_durable_output_is_sequence_zero() {
        assert_eq!(reconcile_durable_output(None, &[], 0, true).unwrap(), 0);
        assert_eq!(reconcile_durable_output(Some(0), &[], 0, true).unwrap(), 0);
    }

    #[test]
    fn hydrate_installs_results_outbox_and_sequence() {
        let mut state = QueryOutputState::new(10);
        let mut results = im::HashMap::new();
        results.insert(1, serde_json::json!({"id": "p1"}));

        let outbox = vec![
            Arc::new(make_query_result("q1", vec![])),
            Arc::new(make_query_result("q1", vec![])),
        ];
        // Sequences on the raw QueryResults are 0 from make_query_result; set them.
        let outbox: Vec<Arc<QueryResult>> = outbox
            .into_iter()
            .enumerate()
            .map(|(i, result)| {
                let mut owned = (*result).clone();
                owned.sequence = (i as u64) + 1;
                Arc::new(owned)
            })
            .collect();

        state.hydrate(results.clone(), outbox.clone(), 2, 0);

        assert_eq!(state.as_of_sequence(), 2);
        assert_eq!(state.results_len(), 1);
        assert_eq!(state.outbox_len(), 2);
        assert_eq!(state.outbox_earliest_seq(), Some(1));
        let fetched = state.fetch_outbox_after(0).unwrap();
        assert_eq!(fetched.len(), 2);
        assert_eq!(fetched[0].sequence, 1);
        assert_eq!(fetched[1].sequence, 2);
    }

    #[test]
    fn hydrate_from_empty_installs_sequence() {
        let mut state = QueryOutputState::new(10);
        assert_eq!(state.as_of_sequence(), 0);
        state.hydrate(im::HashMap::new(), Vec::new(), 4, 0);
        assert_eq!(state.as_of_sequence(), 4);
    }

    #[test]
    fn hydrate_trims_outbox_to_capacity_keeping_newest() {
        let mut state = QueryOutputState::new(2);
        let outbox: Vec<Arc<QueryResult>> = (1..=4)
            .map(|seq| {
                let mut result = make_query_result("q1", vec![]);
                result.sequence = seq;
                Arc::new(result)
            })
            .collect();

        state.hydrate(im::HashMap::new(), outbox, 4, 0);
        assert_eq!(state.outbox_len(), 2);
        assert_eq!(state.outbox_earliest_seq(), Some(3));
        assert_eq!(state.as_of_sequence(), 4);
        let fetched = state.fetch_outbox_after(2).unwrap();
        assert_eq!(
            fetched.iter().map(|r| r.sequence).collect::<Vec<_>>(),
            vec![3, 4]
        );
    }

    #[test]
    fn reset_clears_hydrated_state() {
        let mut state = QueryOutputState::new(10);
        let mut results = im::HashMap::new();
        results.insert(1, serde_json::json!({"id": "p1"}));
        let mut result = make_query_result("q1", vec![]);
        result.sequence = 1;
        state.hydrate(results, vec![Arc::new(result)], 1, 0);
        state.reset();
        assert_eq!(state.as_of_sequence(), 0);
        assert_eq!(state.results_len(), 0);
        assert_eq!(state.outbox_len(), 0);
        assert_eq!(state.generation(), 1);
        assert!(state.initialized());
    }

    #[test]
    fn reset_from_persisted_generation_does_not_restart_at_one() {
        let mut state = QueryOutputState::new(10);
        state.reset_from_generation(4);
        assert_eq!(state.generation(), 5);
        assert!(state.initialized());
        assert_eq!(state.as_of_sequence(), 0);
    }

    #[test]
    fn next_output_generation_always_bumps_from_disk() {
        assert_eq!(next_output_generation(0, 0), 1);
        assert_eq!(next_output_generation(1, 0), 2);
        assert_eq!(next_output_generation(1, 1), 2);
        assert_eq!(next_output_generation(4, 1), 5);
    }

    #[test]
    fn reset_from_generation_bumps_existing_ram_generation() {
        let mut state = QueryOutputState::new(10);
        state.reset_from_generation(1);
        assert_eq!(state.generation(), 2);
        // Resume after a crash that never wrote the new generation still
        // has disk=1. Passing that value must not rewrite generation 1.
        state.reset_from_generation(1);
        assert_eq!(state.generation(), 3);
    }

    #[test]
    fn output_epoch_hash_generation_zero_is_identity() {
        assert_eq!(output_epoch_hash(42, 0), 42);
        assert_ne!(output_epoch_hash(42, 1), 42);
    }
}