haematite 0.6.2

use std::fmt;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use beamr::scheduler::Scheduler;

use beamr::atom::Atom;

use crate::shard::actor::ShardHandle;
use crate::shard::router::{MaterialiseError, ShardRouter};
use crate::sync::endpoint::{DistributionEndpoint, InboundSync};
use crate::sync::protocol::SyncMessage;
use crate::sync::scheduler::SyncSchedulerHandle;

mod config;
mod error;
pub(crate) mod helpers;
mod lock;
mod observer;
mod owner_stamp;
mod receiver;
mod startup;
mod test_support;
pub mod vacuum;

pub use config::{DatabaseConfig, DistributedDatabaseConfig, ON_DISK_FORMAT_VERSION};
pub use error::DatabaseError;
pub use observer::{ObservedEntry, ObserverError, ReadOnlyDatabase};
pub use receiver::respond_to_inbound_writes;

use config::{read_config, validate_database_config};

pub use helpers::run_indexed_parallel;
use helpers::{
    event_range_end, event_range_start, event_sequence_key, has_live_events_on_handle,
    map_shard_error, range_on_handle,
};

const CONFIG_FILE: &str = "config.json";

pub(crate) type DbEntry = (Vec<u8>, Vec<u8>);
pub(crate) type DbRange = Vec<DbEntry>;

/// Top-level database handle. Callers use this API instead of shard actors.
pub struct Database {
    config: DatabaseConfig,
    scheduler: Arc<Scheduler>,
    router: ShardRouter,
    sync_schedulers: Vec<SyncSchedulerHandle>,
    /// Live beamr distribution endpoint, present once `with_distribution` runs.
    ///
    /// The active-active "2a-0" substrate: the inbound drain + outbound send
    /// plumbing that lets two live databases exchange `SyncMessage`s. It does not
    /// yet drive the merge/pull protocol (the sync trigger is still a no-op); it
    /// only exposes the transport primitives later increments build on.
    distribution: Option<DistributionEndpoint>,
    /// AA-3-4a R-LE / R-SEQ: per-shard IN-MEMORY serve-authority. `live_epoch` is
    /// set ONLY by a successful `acquire_shard` THIS lifetime (never recovered
    /// from disk), and the atomic `seq` is drawn once per committed write. The
    /// commit stamp `(live_epoch, seq)` is stamped here on the owner and carried
    /// on the `WriteProposal` so every replica stores the identical stamp.
    owner_stamps: owner_stamp::OwnerStamps,
    /// A4: the exclusive cross-process writer lock on the data dir, held for
    /// this handle's lifetime. Field order keeps it dropped LAST (after shard
    /// teardown in [`Drop`]); the OS also releases it on process death, so a
    /// crashed writer never leaves a stale lock.
    lock: lock::DataDirLock,
    timeout: Duration,
}

impl fmt::Debug for Database {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Database")
            .field("config", &self.config)
            .field("lock", &self.lock)
            .field("timeout", &self.timeout)
            .finish_non_exhaustive()
    }
}

impl Database {
    /// Create a new database directory, write its config, and spawn all shards.
    ///
    /// Acquires the exclusive data-dir writer lock (A4) before writing config;
    /// a live writer already holding the dir fails this loudly with
    /// [`DatabaseError::DataDirLocked`]. A directory that already holds a
    /// database refuses with [`DatabaseError::DataDirAlreadyInitialised`]
    /// (or [`DatabaseError::FormatVersionTooNew`] if a newer binary wrote it)
    /// instead of clobbering its config (A5). A failed create removes the
    /// directory only if this call created it, and does so while still holding
    /// the writer lock — it can never delete a concurrent writer's live dir.
    pub fn create(config: DatabaseConfig) -> Result<Self, DatabaseError> {
        validate_database_config(&config)?;
        startup::initialise_database(config)
    }

    /// Open an existing database directory and restart its shard actors.
    ///
    /// Acquires the exclusive data-dir writer lock (A4): a second writer fails
    /// loudly with [`DatabaseError::DataDirLocked`] instead of corrupting the
    /// shard WALs. For lock-free observation alongside a live writer use
    /// [`ReadOnlyDatabase::open`].
    pub fn open(path: impl AsRef<Path>) -> Result<Self, DatabaseError> {
        let path = path.as_ref().to_path_buf();
        let mut config = read_config(&path)?;
        config.data_dir = path;
        validate_database_config(&config)?;
        startup::open_database(config)
    }

    /// Return the shard index that owns `key`.
    pub fn shard_for(&self, key: &[u8]) -> usize {
        self.router.shard_for(key)
    }

    /// Atomically append event entries under `key` using optimistic concurrency.
    pub fn append(
        &self,
        key: Vec<u8>,
        entries: Vec<Vec<u8>>,
        expected_seq: u64,
    ) -> Result<u64, DatabaseError> {
        self.append_with_ttl(key, entries, expected_seq, None)
    }

    /// Atomically append event entries with optional TTL metadata.
    pub fn append_with_ttl(
        &self,
        key: Vec<u8>,
        entries: Vec<Vec<u8>>,
        expected_seq: u64,
        ttl: Option<Duration>,
    ) -> Result<u64, DatabaseError> {
        self.handle_for(&key)?
            .append_with_ttl(key, entries, expected_seq, ttl, self.timeout)
            .map_err(map_shard_error)
    }

    /// Read all appended events for `key` in sequence order.
    pub fn read_events(&self, key: &[u8]) -> Result<Vec<Vec<u8>>, DatabaseError> {
        self.read_events_from(key, 0)
    }

    /// Read appended events for `key` from `from_seq` onward.
    pub fn read_events_from(
        &self,
        key: &[u8],
        from_seq: u64,
    ) -> Result<Vec<Vec<u8>>, DatabaseError> {
        let from = event_range_start(key, from_seq);
        let to = event_range_end(key);
        let handle = self.handle_for(key)?;
        let entries = range_on_handle(&handle, &from, &to, self.timeout)?;
        Ok(entries.into_iter().map(|(_, value)| value).collect())
    }

    /// Read appended event entries for `key` from `from_seq` onward as raw
    /// `(encoded_key, value)` pairs, in sequence order.
    ///
    /// Unlike [`Self::read_events_from`], this preserves the encoded tree key so
    /// the caller (the `EventStore`) can decode each event's sequence number from
    /// its key rather than trusting a value-side copy.
    pub fn read_event_entries_from(
        &self,
        key: &[u8],
        from_seq: u64,
    ) -> Result<DbRange, DatabaseError> {
        let from = event_range_start(key, from_seq);
        let to = event_range_end(key);
        let handle = self.handle_for(key)?;
        range_on_handle(&handle, &from, &to, self.timeout)
    }

    /// Read the next sequence metadata for an event stream, if the stream exists.
    pub fn read_stream_next_seq(&self, key: &[u8]) -> Result<Option<u64>, DatabaseError> {
        // Route on the STREAM key — `append` writes the sequence metadata into the
        // shard of the stream key, so the read must select the same shard. Routing
        // on `event_sequence_key(key)` (a different hash) would read the wrong
        // shard for `shard_count > 1` and miss the metadata entirely.
        self.handle_for(key)?
            .read_value(event_sequence_key(key), self.timeout)
            .map_err(map_shard_error)
    }

    /// Return true if a stream has at least one non-expired event visible now.
    pub fn stream_has_live_events(&self, key: &[u8]) -> Result<bool, DatabaseError> {
        let handle = self.handle_for(key)?;
        has_live_events_on_handle(&handle, key, self.timeout)
    }

    /// Read the scalar `u64` value at `key`, or `None` if it is unset.
    pub fn read_value(&self, key: &[u8]) -> Result<Option<u64>, DatabaseError> {
        self.handle_for(key)?
            .read_value(key.to_vec(), self.timeout)
            .map_err(map_shard_error)
    }

    /// Atomically compare-and-swap the scalar `u64` value at `key`.
    ///
    /// The read-compare-write executes inside the owning shard's single-threaded
    /// actor, so concurrent CAS calls against the same key are serialised and
    /// cannot race. Returns [`DatabaseError::CasMismatch`] if the current value
    /// is not `expected`.
    pub fn cas(&self, key: Vec<u8>, expected: Option<u64>, new: u64) -> Result<(), DatabaseError> {
        self.handle_for(&key)?
            .cas(key, expected, new, self.timeout)
            .map_err(map_shard_error)
    }

    /// Write the single-node GENESIS `cluster/members` record (CSOT-1, task #146).
    ///
    /// A lone node writes its own denominator-1, self-quorum member set at config
    /// epoch 0 (#146 §4.2 step 2) and can immediately read it back with
    /// [`Self::read_cluster_members`]. The record is persisted under the reserved
    /// [`crate::sync::CLUSTER_MEMBERS_KEY`] via the existing durable append
    /// primitive; genesis is the append at sequence 0, so a second genesis attempt
    /// on an already-formed cluster fails with a sequence conflict rather than
    /// silently overwriting the durable record.
    ///
    /// This is INERT: it does not change any quorum/send behaviour by itself.
    /// `resolve_membership` only consults the record once a caller reads it and
    /// passes it to [`crate::sync::resolve_membership_with_record`]. CSOT-1 writes
    /// no deltas (join/leave/evict are later phases).
    pub fn write_genesis_cluster_members(
        &self,
        record: &crate::sync::ClusterMembers,
    ) -> Result<(), DatabaseError> {
        let bytes = record.encode()?;
        self.append(crate::sync::CLUSTER_MEMBERS_KEY.to_vec(), vec![bytes], 0)?;
        Ok(())
    }

    /// Read the durable `cluster/members` record, or `None` if no record exists yet
    /// (a fresh, never-formed cluster) (CSOT-1, task #146).
    ///
    /// `None` is the load-bearing FALLBACK signal: with no durable record,
    /// `resolve_membership_with_record(config, None, ..)` sizes quorum from static
    /// `config.nodes`, byte-identical to the pre-CSOT-1 path. When a record IS
    /// present, the LATEST stored version is returned (later deltas append newer
    /// versions; the newest wins) and its denominator takes precedence.
    pub fn read_cluster_members(
        &self,
    ) -> Result<Option<crate::sync::ClusterMembers>, DatabaseError> {
        let versions = self.read_events(crate::sync::CLUSTER_MEMBERS_KEY)?;
        let Some(latest) = versions.last() else {
            return Ok(None);
        };
        let record = crate::sync::ClusterMembers::decode(latest)?;
        Ok(Some(record))
    }

    /// Attach a live beamr distribution endpoint to this database.
    ///
    /// This is the active-active "2a-0" substrate: it installs the inbound-drain
    /// and outbound-send plumbing two live databases need to exchange
    /// `SyncMessage`s over a real network. The endpoint owns its own atom table,
    /// connection manager, accept loop, and tokio runtime (see
    /// [`DistributionEndpoint`]).
    ///
    /// It does NOT replace the no-op sync trigger or drive the pull/merge
    /// protocol — those are later increments. The database simply takes ownership
    /// of the endpoint and re-exports its transport primitives
    /// ([`Database::connect_peer`], [`Database::send_sync_message`],
    /// [`Database::recv_sync_message`]).
    #[must_use]
    pub fn with_distribution(mut self, endpoint: DistributionEndpoint) -> Self {
        self.distribution = Some(endpoint);
        self
    }

    /// Borrow the attached distribution endpoint, if any.
    #[must_use]
    pub const fn distribution(&self) -> Option<&DistributionEndpoint> {
        self.distribution.as_ref()
    }

    /// Register `peer_name` at `addr` and dial it over real distribution.
    ///
    /// Requires [`Database::with_distribution`] to have installed an endpoint.
    pub fn connect_peer(
        &self,
        peer_name: &str,
        addr: std::net::SocketAddr,
    ) -> Result<(), DatabaseError> {
        let endpoint = self.require_distribution()?;
        endpoint.add_peer(peer_name, addr);
        endpoint
            .connect(peer_name)
            .map_err(|error| DatabaseError::Distribution(error.to_string()))
    }

    /// Intern `peer_name` into the endpoint's atom table for addressed sends.
    pub fn peer_atom(&self, peer_name: &str) -> Result<Atom, DatabaseError> {
        Ok(self.require_distribution()?.peer_atom(peer_name))
    }

    /// Send `message` to the peer named `peer_name` over the live transport.
    ///
    /// Requires an attached endpoint and an established connection to the peer.
    pub fn send_sync_message(
        &self,
        peer_name: &str,
        message: &SyncMessage,
    ) -> Result<(), DatabaseError> {
        self.require_distribution()?
            .send_to(peer_name, message)
            .map_err(|error| DatabaseError::Distribution(error.to_string()))
    }

    /// Block until an inbound sync message arrives or `timeout` elapses.
    ///
    /// Returns `Ok(Some(_))` with the decoded message (or a decode error from the
    /// wire), `Ok(None)` on timeout, and an error if no endpoint is attached or
    /// the drain has been disconnected.
    pub fn recv_sync_message(
        &self,
        timeout: Duration,
    ) -> Result<Option<InboundSync>, DatabaseError> {
        self.require_distribution()?
            .recv_inbound(timeout)
            .map_err(|error| DatabaseError::Distribution(error.to_string()))
    }

    /// Atoms for all currently active distribution connections.
    pub fn connected_nodes(&self) -> Result<Vec<Atom>, DatabaseError> {
        Ok(self.require_distribution()?.connected_nodes())
    }

    fn require_distribution(&self) -> Result<&DistributionEndpoint, DatabaseError> {
        self.distribution
            .as_ref()
            .ok_or_else(|| DatabaseError::Distribution("no distribution endpoint".into()))
    }

    /// Collect every stream's `(stream_key, next_seq)` pair across all shards.
    ///
    /// This walks each shard in parallel, scanning its full key range for the
    /// per-stream sequence-metadata keys and decoding each one. It is the
    /// O(total entries) traversal that backs the `EventStore` `scan` predicate.
    pub fn scan_sequence_keys(&self) -> Result<Vec<(Vec<u8>, u64)>, DatabaseError> {
        // Only MATERIALISED shards can hold streams: an un-materialised shard has
        // never had a write, so it contributes no sequence keys. Scanning the
        // materialised set (not `0..shard_count`) is both correct and O(used).
        let handles = self.router.materialised_handles();
        let timeout = self.timeout;
        let results = run_indexed_parallel(handles, |handle: ShardHandle| {
            handle.scan_sequences(timeout)
        })?;
        let mut streams = Vec::new();
        for (_, result) in results {
            streams.extend(result.map_err(map_shard_error)?);
        }
        Ok(streams)
    }

    /// Collect `(stream_key, next_seq)` pairs from ONLY the named shards.
    ///
    /// The scoped counterpart of [`Self::scan_sequence_keys`]: a node that owns a
    /// subset of shards enumerates only its own streams (e.g. to recover exactly
    /// the workflows whose event streams it serves) without paying for, or
    /// surfacing, streams that live on shards another node owns. Each id must be in
    /// `0..shard_count`; an out-of-range id is [`DatabaseError::InvalidShardCount`].
    pub fn scan_sequence_keys_for_shards(
        &self,
        shard_ids: &[usize],
    ) -> Result<Vec<(Vec<u8>, u64)>, DatabaseError> {
        // Materialise each requested shard on demand: a node recovering exactly
        // the streams for shards it adopts must open (and WAL-recover) those
        // shards even if they were never touched this lifetime. An out-of-range id
        // is still rejected as `InvalidShardCount`.
        let mut handles = Vec::with_capacity(shard_ids.len());
        for &shard_id in shard_ids {
            handles.push(self.handle_for_shard(shard_id)?);
        }
        let timeout = self.timeout;
        let results = run_indexed_parallel(handles, |handle: ShardHandle| {
            handle.scan_sequences(timeout)
        })?;
        let mut streams = Vec::new();
        for (_, result) in results {
            streams.extend(result.map_err(map_shard_error)?);
        }
        Ok(streams)
    }

    pub const fn shard_count(&self) -> usize {
        self.config.shard_count
    }

    /// Map a lazy-materialisation failure into the public [`DatabaseError`].
    ///
    /// A first-touch spawn failure (bad directory, scheduler refusal) surfaces as
    /// a [`DatabaseError::ShardSpawn`]; an out-of-range shard id surfaces as
    /// [`DatabaseError::InvalidShardCount`] to preserve the old routing contract.
    fn map_materialise_error(&self, error: MaterialiseError) -> DatabaseError {
        if error.shard_id >= self.config.shard_count {
            DatabaseError::InvalidShardCount
        } else {
            DatabaseError::ShardSpawn(error.message)
        }
    }

    pub(crate) const fn timeout(&self) -> Duration {
        self.timeout
    }

    /// A single consistent snapshot of the materialised shards: their ids and
    /// their handles, index-aligned (`ids[i]` owns `handles[i]`). Used by the
    /// commit path to map each parallel result back to its real shard id before
    /// synthesising the empty root for the un-materialised slots (GATE 1).
    pub(crate) fn materialised_shards(&self) -> (Vec<usize>, Vec<ShardHandle>) {
        self.router.materialised_snapshot()
    }

    /// The shard ids materialised so far, ascending. Test-only projection used by
    /// the lazy-root spike to assert exactly which shards a lazy workload touched
    /// (and that a force-materialised DB touched every shard).
    #[cfg(test)]
    pub(crate) fn materialised_shard_ids(&self) -> Vec<usize> {
        self.router.materialised_shard_ids()
    }

    /// Materialise-on-miss the shard owning `key` and return a handle clone.
    pub(crate) fn handle_for(&self, key: &[u8]) -> Result<ShardHandle, DatabaseError> {
        self.router
            .handle_for(key)
            .map_err(|error| self.map_materialise_error(error))
    }

    /// Route to a shard by its index (AA-3-2 election routing). A `Prepare`/
    /// `acquire_shard` names the shard directly, so it bypasses key-hash routing.
    /// Materialisation runs the shard's normal boot (store open + durable WAL/
    /// promise recovery) before the handle serves — GATE 3.
    pub(crate) fn handle_for_shard(&self, shard_id: usize) -> Result<ShardHandle, DatabaseError> {
        self.router
            .handle_for_shard(shard_id)
            .map_err(|error| self.map_materialise_error(error))
    }
}

impl Drop for Database {
    fn drop(&mut self) {
        for handle in &self.sync_schedulers {
            if let Err(error) = handle.shutdown(self.timeout) {
                log::debug!(
                    "database sync scheduler shutdown skipped for supervisor pid {}: {error}",
                    handle.supervisor_pid()
                );
            }
        }
        // Sweeps and shard actors both live in the router now (materialised
        // together, torn down together).
        self.router.shutdown_all(self.timeout);
        self.scheduler.shutdown();
    }
}

#[cfg(test)]
#[path = "db_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "db/lazy_root_spike_tests.rs"]
mod lazy_root_spike_tests;

#[cfg(test)]
#[path = "db/format_version_tests.rs"]
mod format_version_tests;

#[cfg(test)]
#[path = "db/ownership_transfer_tests.rs"]
mod ownership_transfer_tests;