pub mod disk;
pub mod fs;
pub mod replay;

pub use fs::FsStorage;

use std::future::Future;
use std::path::PathBuf;

use crate::error::Result;
use crate::types::{NamespaceId, PublishedMessage, TopicName, TopicOffset};

/// Trait for durable message storage backends.
///
/// Implementors must be safe to share across tasks (`Send + Sync`).
/// All async methods must return `Send` futures so the engine can use
/// them from `tokio::spawn` contexts.
///
/// [`append`](Self::append) writes to the OS page cache but does not fsync.
/// Callers must call [`flush`](Self::flush) for per-topic durability or
/// [`shutdown`](Self::shutdown) to flush all topics before dropping.
pub trait StorageBackend: Send + Sync {
    /// Initialize the storage backend (e.g., create directories, pre-populate caches).
    fn init(&self) -> impl Future<Output = Result<()>> + Send;

    /// Append a published message to the log for the given namespace/topic.
    ///
    /// Writes to the OS page cache but does not fsync. Call [`flush`](Self::flush)
    /// to ensure durability.
    fn append(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        msg: &PublishedMessage,
    ) -> impl Future<Output = Result<()>> + Send;

    /// Fsync the active segment for the given namespace/topic.
    ///
    /// No-op if the topic has no active writer.
    fn flush(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
    ) -> impl Future<Output = Result<()>> + Send;

    /// Replay all persisted messages for a given namespace/topic, in offset order.
    fn replay(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
    ) -> impl Future<Output = Result<Vec<PublishedMessage>>> + Send;

    /// Enumerate all known namespace/topic pairs that have persisted data.
    ///
    /// Used by the queue engine on startup to discover which topics need replay.
    fn list_topics(&self) -> impl Future<Output = Result<Vec<(NamespaceId, TopicName)>>> + Send;

    /// Return the directory path where segments are stored for this topic.
    fn segments_dir(&self, namespace: &NamespaceId, topic: &TopicName) -> PathBuf;

    /// Read a single message by offset from persisted segments.
    ///
    /// Returns `Ok(None)` if the offset is not found on disk (deleted by
    /// segment rotation or never persisted). This is the cold-path
    /// counterpart to the in-memory `Topic::consume_from`.
    fn read_at(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        offset: TopicOffset,
    ) -> impl Future<Output = Result<Option<PublishedMessage>>> + Send;

    /// Read up to `max_count` messages starting from `start_offset` from
    /// persisted segments. Returns messages in offset order.
    fn read_cold_batch(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        start_offset: TopicOffset,
        max_count: usize,
    ) -> impl Future<Output = Result<Vec<PublishedMessage>>> + Send;

    /// Delete segment files matching the given criteria. Returns count deleted.
    ///
    /// - `before_offset`: delete segments with base_offset < this value
    /// - `max_segments`: keep at most this many segments (delete oldest excess)
    /// - `safe_below`: never delete a segment whose base_offset >= this floor
    ///
    /// All filters compose: a segment must match all applicable criteria to be
    /// deleted. Pass `None` to disable a criterion.
    fn delete_segments(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        before_offset: Option<TopicOffset>,
        max_segments: Option<usize>,
        safe_below: Option<TopicOffset>,
    ) -> impl Future<Output = Result<usize>> + Send;

    /// Find the offset of the first message at or after `timestamp_millis`.
    /// Returns `None` if no message exists at or after the timestamp.
    fn offset_at_timestamp(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        timestamp_millis: u64,
    ) -> impl Future<Output = Result<Option<TopicOffset>>> + Send;

    /// Merge runs of consecutive small segments for a topic.
    /// Returns the number of merge operations performed (each merging 2+ segments).
    /// `safe_below`: never merge segments whose base_offset >= this floor
    /// (active consumers may have open cursors on those segments).
    fn compact_segments(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
        merge_threshold_bytes: u64,
        safe_below: Option<TopicOffset>,
    ) -> impl Future<Output = Result<usize>> + Send;

    /// Return the total bytes on disk for a namespace's segments.
    fn namespace_disk_bytes(&self, namespace: &NamespaceId) -> impl Future<Output = u64> + Send;

    /// Flush all topics that have pending writes.
    ///
    /// Called before entering drain mode to ensure durability of all buffered
    /// data before clients disconnect.
    fn flush_all(&self) -> impl Future<Output = Result<()>> + Send;

    /// Flush and close the active writer for a topic, evicting it from the fd cache.
    /// No-op if no active writer exists. Must be called before `delete_topic_dir`.
    fn close_topic_writer(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
    ) -> impl Future<Output = Result<()>> + Send;

    /// Delete the entire segment directory for a topic.
    /// No-op if the directory does not exist.
    fn delete_topic_dir(
        &self,
        namespace: &NamespaceId,
        topic: &TopicName,
    ) -> impl Future<Output = Result<()>> + Send;

    /// Flush and close all active writers. Call before dropping.
    fn shutdown(&self) -> impl Future<Output = Result<()>> + Send;
}

#[cfg(any(test, feature = "test-support"))]
pub mod mock {
    use super::*;
    use std::collections::HashMap;
    use std::sync::atomic::{AtomicU64, Ordering};
    use tokio::sync::Mutex;

    /// In-memory mock storage for testing downstream modules.
    pub struct MockStorage {
        data: Mutex<HashMap<(String, String), Vec<PublishedMessage>>>,
        namespace_bytes: Mutex<HashMap<String, AtomicU64>>,
    }

    impl MockStorage {
        pub fn new() -> Self {
            Self {
                data: Mutex::new(HashMap::new()),
                namespace_bytes: Mutex::new(HashMap::new()),
            }
        }
    }

    impl StorageBackend for MockStorage {
        async fn init(&self) -> Result<()> {
            Ok(())
        }

        async fn append(
            &self,
            namespace: &NamespaceId,
            topic: &TopicName,
            msg: &PublishedMessage,
        ) -> Result<()> {
            // Track the full encoded record size (same as FsStorage) for
            // consistent disk quota behavior between mock and real storage.
            let record_bytes = crate::storage::fs::encode_record(msg)
                .map(|r| r.len() as u64)
                .unwrap_or(msg.envelope.payload.len() as u64);
            let key = (namespace.0.clone(), topic.0.clone());
            let mut data = self.data.lock().await;
            data.entry(key).or_default().push(msg.clone());
            let mut bytes = self.namespace_bytes.lock().await;
            bytes
                .entry(namespace.0.clone())
                .or_insert_with(|| AtomicU64::new(0))
                .fetch_add(record_bytes, Ordering::Relaxed);
            Ok(())
        }

        async fn flush(&self, _namespace: &NamespaceId, _topic: &TopicName) -> Result<()> {
            Ok(())
        }

        async fn replay(
            &self,
            namespace: &NamespaceId,
            topic: &TopicName,
        ) -> Result<Vec<PublishedMessage>> {
            let key = (namespace.0.clone(), topic.0.clone());
            let data = self.data.lock().await;
            Ok(data.get(&key).cloned().unwrap_or_default())
        }

        async fn list_topics(&self) -> Result<Vec<(NamespaceId, TopicName)>> {
            let data = self.data.lock().await;
            let topics = data
                .keys()
                .map(|(ns, topic)| (NamespaceId(ns.clone()), TopicName(topic.clone())))
                .collect();
            Ok(topics)
        }

        fn segments_dir(&self, namespace: &NamespaceId, topic: &TopicName) -> PathBuf {
            PathBuf::from(format!("mock/{}/{}/segments", namespace.0, topic.0))
        }

        async fn read_at(
            &self,
            namespace: &NamespaceId,
            topic: &TopicName,
            offset: TopicOffset,
        ) -> Result<Option<PublishedMessage>> {
            let key = (namespace.0.clone(), topic.0.clone());
            let data = self.data.lock().await;
            Ok(data
                .get(&key)
                .and_then(|msgs| msgs.iter().find(|m| m.offset == offset).cloned()))
        }

        async fn read_cold_batch(
            &self,
            namespace: &NamespaceId,
            topic: &TopicName,
            start_offset: TopicOffset,
            max_count: usize,
        ) -> Result<Vec<PublishedMessage>> {
            let key = (namespace.0.clone(), topic.0.clone());
            let data = self.data.lock().await;
            Ok(data
                .get(&key)
                .map(|msgs| {
                    msgs.iter()
                        .filter(|m| m.offset >= start_offset)
                        .take(max_count)
                        .cloned()
                        .collect()
                })
                .unwrap_or_default())
        }

        async fn delete_segments(
            &self,
            _namespace: &NamespaceId,
            _topic: &TopicName,
            _before_offset: Option<TopicOffset>,
            _max_segments: Option<usize>,
            _safe_below: Option<TopicOffset>,
        ) -> Result<usize> {
            Ok(0)
        }

        async fn compact_segments(
            &self,
            _namespace: &NamespaceId,
            _topic: &TopicName,
            _merge_threshold_bytes: u64,
            _safe_below: Option<TopicOffset>,
        ) -> Result<usize> {
            Ok(0)
        }

        async fn offset_at_timestamp(
            &self,
            namespace: &NamespaceId,
            topic: &TopicName,
            timestamp_millis: u64,
        ) -> Result<Option<TopicOffset>> {
            let key = (namespace.0.clone(), topic.0.clone());
            let data = self.data.lock().await;
            Ok(data
                .get(&key)
                .and_then(|msgs| {
                    msgs.iter()
                        .find(|m| m.envelope.timestamp >= timestamp_millis)
                })
                .map(|m| m.offset))
        }

        async fn namespace_disk_bytes(&self, namespace: &NamespaceId) -> u64 {
            let bytes = self.namespace_bytes.lock().await;
            bytes
                .get(&namespace.0)
                .map(|v| v.load(Ordering::Relaxed))
                .unwrap_or(0)
        }

        async fn flush_all(&self) -> Result<()> {
            Ok(())
        }

        async fn close_topic_writer(&self, _ns: &NamespaceId, _topic: &TopicName) -> Result<()> {
            Ok(())
        }

        async fn delete_topic_dir(&self, _ns: &NamespaceId, _topic: &TopicName) -> Result<()> {
            Ok(())
        }

        async fn shutdown(&self) -> Result<()> {
            Ok(())
        }
    }
}
