//! Storage primitives used by `chainfire-raft`.
//!
//! In production (`rocksdb-storage` feature), we re-export the real ChainFire storage layer.
//! For lightweight testing/simulation (default), we provide a small in-memory implementation
//! that avoids native dependencies (RocksDB/libclang).

#[cfg(feature = "rocksdb-storage")]
pub use chainfire_storage::{
    EntryPayload, LogEntry, LogId, LogState, LogStorage, StateMachine, Vote,
};

#[cfg(not(feature = "rocksdb-storage"))]
mod mem {
    use chainfire_types::command::{RaftCommand, RaftResponse};
    use chainfire_types::error::StorageError;
    use chainfire_types::kv::{KvEntry, Revision};
    use parking_lot::RwLock;
    use serde::{Deserialize, Serialize};
    use std::collections::{BTreeMap, HashMap};
    use std::ops::RangeBounds;
    use std::sync::atomic::{AtomicU64, Ordering};

    pub type LogIndex = u64;
    pub type Term = u64;

    /// Log ID combining term and index.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default)]
    pub struct LogId {
        pub term: Term,
        pub index: LogIndex,
    }

    /// Payload of a log entry.
    #[derive(Debug, Clone, Serialize, Deserialize)]
    pub enum EntryPayload<D> {
        /// A blank entry for leader establishment.
        Blank,
        /// A normal data entry.
        Normal(D),
        /// Membership change entry.
        Membership(Vec<u64>),
    }

    /// A log entry stored in the Raft log.
    #[derive(Debug, Clone, Serialize, Deserialize)]
    pub struct LogEntry<D> {
        pub log_id: LogId,
        pub payload: EntryPayload<D>,
    }

    /// Persisted vote information.
    #[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
    pub struct Vote {
        pub term: Term,
        pub node_id: Option<u64>,
        pub committed: bool,
    }

    /// Log storage state.
    #[derive(Debug, Clone, Default)]
    pub struct LogState {
        pub last_purged_log_id: Option<LogId>,
        pub last_log_id: Option<LogId>,
    }

    /// In-memory Raft log storage.
    ///
    /// Stores bincode-encoded `LogEntry<D>` blobs keyed by log index.
    pub struct LogStorage {
        vote: RwLock<Option<Vote>>,
        logs: RwLock<BTreeMap<LogIndex, Vec<u8>>>,
        last_purged_log_id: RwLock<Option<LogId>>,
    }

    impl Default for LogStorage {
        fn default() -> Self {
            Self::new_in_memory()
        }
    }

    impl LogStorage {
        pub fn new_in_memory() -> Self {
            Self {
                vote: RwLock::new(None),
                logs: RwLock::new(BTreeMap::new()),
                last_purged_log_id: RwLock::new(None),
            }
        }

        pub fn store(&self) -> chainfire_storage::RocksStore {
            // This is a hack to satisfy the API. In memory mode, we shouldn't really
            // be calling this if we want to avoid RocksDB, but chainfire-api expects it.
            panic!("LogStorage::store() called in memory mode");
        }

        pub fn get_log_state(&self) -> Result<LogState, StorageError> {
            let last_purged_log_id = *self.last_purged_log_id.read();
            let logs = self.logs.read();
            let last_log_id = match logs.iter().next_back() {
                Some((_idx, bytes)) if !bytes.is_empty() => {
                    match bincode::deserialize::<LogEntry<Vec<u8>>>(bytes) {
                        Ok(entry) => Some(entry.log_id),
                        Err(e) => {
                            eprintln!(
                                "Warning: Failed to deserialize log entry in mem storage: {e}, treating as empty log"
                            );
                            last_purged_log_id
                        }
                    }
                }
                _ => last_purged_log_id,
            };

            Ok(LogState {
                last_purged_log_id,
                last_log_id,
            })
        }

        pub fn save_vote(&self, vote: Vote) -> Result<(), StorageError> {
            *self.vote.write() = Some(vote);
            Ok(())
        }

        pub fn read_vote(&self) -> Result<Option<Vote>, StorageError> {
            Ok(*self.vote.read())
        }

        pub fn append<D: Serialize>(&self, entries: &[LogEntry<D>]) -> Result<(), StorageError> {
            if entries.is_empty() {
                return Ok(());
            }
            let mut logs = self.logs.write();
            for entry in entries {
                let bytes = bincode::serialize(entry)
                    .map_err(|e| StorageError::Serialization(e.to_string()))?;
                logs.insert(entry.log_id.index, bytes);
            }
            Ok(())
        }

        pub fn get_log_entries<D: for<'de> Deserialize<'de>>(
            &self,
            range: impl RangeBounds<LogIndex>,
        ) -> Result<Vec<LogEntry<D>>, StorageError> {
            let logs = self.logs.read();

            let start = match range.start_bound() {
                std::ops::Bound::Included(&idx) => idx,
                std::ops::Bound::Excluded(&idx) => idx + 1,
                std::ops::Bound::Unbounded => 0,
            };

            let end = match range.end_bound() {
                std::ops::Bound::Included(&idx) => Some(idx),
                std::ops::Bound::Excluded(&idx) => Some(idx.saturating_sub(1)),
                std::ops::Bound::Unbounded => None,
            };

            let iter: Box<dyn Iterator<Item = (&LogIndex, &Vec<u8>)> + '_> = match end {
                Some(end_inclusive) => Box::new(logs.range(start..=end_inclusive)),
                None => Box::new(logs.range(start..)),
            };

            let mut out = Vec::new();
            for (_idx, bytes) in iter {
                let entry: LogEntry<D> = bincode::deserialize(bytes)
                    .map_err(|e| StorageError::Serialization(e.to_string()))?;
                out.push(entry);
            }
            Ok(out)
        }

        pub fn truncate(&self, from_index: LogIndex) -> Result<(), StorageError> {
            let mut logs = self.logs.write();
            // Remove all entries >= from_index
            let _ = logs.split_off(&from_index);
            Ok(())
        }

        #[allow(dead_code)]
        pub fn purge_with_log_id(&self, log_id: LogId) -> Result<(), StorageError> {
            // In-memory compaction marker only; entries are not retained once purged.
            *self.last_purged_log_id.write() = Some(log_id);
            self.truncate(log_id.index + 1)?;
            Ok(())
        }
    }

    /// Minimal in-memory KV store used by the in-memory state machine.
    pub struct KvStore {
        data: RwLock<HashMap<Vec<u8>, KvEntry>>,
        revision: AtomicU64,
    }

    impl Default for KvStore {
        fn default() -> Self {
            Self {
                data: RwLock::new(HashMap::new()),
                revision: AtomicU64::new(0),
            }
        }
    }

    impl KvStore {
        pub fn current_revision(&self) -> Revision {
            self.revision.load(Ordering::SeqCst)
        }

        fn next_revision(&self) -> Revision {
            self.revision.fetch_add(1, Ordering::SeqCst) + 1
        }

        pub fn get(&self, key: &[u8]) -> Result<Option<KvEntry>, StorageError> {
            Ok(self.data.read().get(key).cloned())
        }

        pub fn range_count(&self, start: &[u8], end: Option<&[u8]>) -> Result<usize, StorageError> {
            let data = self.data.read();
            let count = if let Some(end) = end {
                data.iter().filter(|(k, _)| k.as_slice() >= start && k.as_slice() < end).count()
            } else {
                data.iter().filter(|(k, _)| k.as_slice() >= start).count()
            };
            Ok(count)
        }

        pub fn range_with_limit(&self, start: &[u8], end: Option<&[u8]>, limit: Option<usize>) -> Result<(Vec<KvEntry>, bool), StorageError> {
            let data = self.data.read();
            let mut entries: Vec<_> = if let Some(end) = end {
                data.iter()
                    .filter(|(k, _)| k.as_slice() >= start && k.as_slice() < end)
                    .map(|(_, v)| v.clone())
                    .collect()
            } else {
                data.iter()
                    .filter(|(k, _)| k.as_slice() >= start)
                    .map(|(_, v)| v.clone())
                    .collect()
            };
            entries.sort_by(|a, b| a.key.cmp(&b.key));
            
            if let Some(limit) = limit {
                let more = entries.len() > limit;
                entries.truncate(limit);
                Ok((entries, more))
            } else {
                Ok((entries, false))
            }
        }

        pub fn set_revision(&self, revision: Revision) {
            self.revision.store(revision, Ordering::SeqCst);
        }

        pub fn put(
            &self,
            key: Vec<u8>,
            value: Vec<u8>,
            lease_id: Option<i64>,
        ) -> Result<(Revision, Option<KvEntry>), StorageError> {
            let mut data = self.data.write();
            let prev = data.get(&key).cloned();
            let revision = self.next_revision();

            let entry = match &prev {
                Some(old) => old.update(value, revision),
                None => {
                    if let Some(lease) = lease_id {
                        KvEntry::with_lease(key.clone(), value, revision, lease)
                    } else {
                        KvEntry::new(key.clone(), value, revision)
                    }
                }
            };
            data.insert(key, entry);
            Ok((revision, prev))
        }

        pub fn delete(&self, key: &[u8]) -> Result<(Revision, Option<KvEntry>), StorageError> {
            let mut data = self.data.write();
            let prev = data.remove(key);
            let revision = self.next_revision();
            Ok((revision, prev))
        }
    }

    /// Minimal in-memory state machine for Raft simulation.
    pub struct StateMachine {
        kv: KvStore,
    }

    pub struct LeaseStore;
    impl LeaseStore {
        pub fn list(&self) -> Vec<chainfire_types::lease::Lease> { vec![] }
    }

    impl Default for StateMachine {
        fn default() -> Self {
            Self::new_in_memory()
        }
    }

    impl StateMachine {
        pub fn new_in_memory() -> Self {
            Self { kv: KvStore::default() }
        }

        pub fn kv(&self) -> &KvStore {
            &self.kv
        }

        pub fn current_revision(&self) -> Revision {
            self.kv.current_revision()
        }

        pub fn leases(&self) -> LeaseStore {
            LeaseStore
        }

        pub fn apply(&self, command: RaftCommand) -> Result<RaftResponse, StorageError> {
            match command {
                RaftCommand::Put {
                    key,
                    value,
                    lease_id,
                    prev_kv,
                } => {
                    let (rev, prev) = self.kv.put(key, value, lease_id)?;
                    Ok(RaftResponse::with_prev_kv(rev, if prev_kv { prev } else { None }))
                }
                RaftCommand::Delete { key, prev_kv } => {
                    let (rev, prev) = self.kv.delete(&key)?;
                    let deleted = if prev.is_some() { 1 } else { 0 };
                    Ok(RaftResponse {
                        revision: rev,
                        prev_kv: if prev_kv { prev } else { None },
                        deleted,
                        ..Default::default()
                    })
                }
                RaftCommand::Noop => Ok(RaftResponse::new(self.current_revision())),
                other => Err(StorageError::Serialization(format!(
                    "mem state machine: unsupported command variant: {other:?}"
                ))),
            }
        }
    }

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

        #[test]
        fn mem_log_storage_append_and_get() {
            let storage = LogStorage::new_in_memory();
            let entries = vec![
                LogEntry {
                    log_id: LogId { term: 1, index: 1 },
                    payload: EntryPayload::Normal(b"a".to_vec()),
                },
                LogEntry {
                    log_id: LogId { term: 1, index: 2 },
                    payload: EntryPayload::Normal(b"b".to_vec()),
                },
            ];
            storage.append(&entries).unwrap();
            let got: Vec<LogEntry<Vec<u8>>> = storage.get_log_entries(1..=2).unwrap();
            assert_eq!(got.len(), 2);
            assert_eq!(got[0].log_id.index, 1);
        }
    }
}

#[cfg(not(feature = "rocksdb-storage"))]
pub use mem::{EntryPayload, LogEntry, LogId, LogState, LogStorage, StateMachine, Vote};


