//! Persistent Raft storage for FlareDB
//!
//! This module provides a RocksDB-backed implementation of OpenRaft's RaftStorage trait.

use crate::storage::{FlareSnapshot, FlareStateMachine};
use crate::types::{
    FlareEntry, FlareLogId, FlareNode, FlareNodeId, FlareResponse, FlareTypeConfig,
};
use flaredb_storage::rocks_engine::RocksEngine;
use openraft::storage::{LogFlushed, LogState, RaftLogReader, RaftLogStorage, RaftSnapshotBuilder, RaftStateMachine, Snapshot};
use openraft::{EntryPayload, OptionalSend, RaftLogId, SnapshotMeta};
use openraft::{StorageError, StorageIOError, StoredMembership, Vote};
use std::fmt::Debug;
use std::io::Cursor;
use std::ops::RangeBounds;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, trace};

// Raft state keys
const KEY_VOTE: &str = "vote";
const KEY_COMMITTED: &str = "committed";
const KEY_LAST_PURGED: &str = "last_purged";
const KEY_SNAPSHOT_DATA: &str = "snapshot_data";
const KEY_SNAPSHOT_META: &str = "snapshot_meta";
const KEY_SNAPSHOT_IDX: &str = "snapshot_idx";

/// Persistent storage implementing RaftStorage for FlareDB
pub struct PersistentFlareStore {
    /// RocksDB engine
    engine: Arc<RocksEngine>,
    /// State machine (still in-memory for now, but log is persistent)
    sm: RwLock<FlareStateMachine>,
    /// Snapshot counter
    snapshot_idx: RwLock<u64>,
    /// Current snapshot
    current_snapshot: RwLock<Option<FlareSnapshot>>,
}

impl PersistentFlareStore {
    /// Create a new PersistentFlareStore with recovery from disk
    ///
    /// On startup, this will:
    /// 1. Load the last snapshot from RocksDB (if any)
    /// 2. Restore the state machine from the snapshot
    /// 3. OpenRaft will replay logs since the snapshot automatically
    pub fn new(engine: Arc<RocksEngine>) -> Self {
        // Try to recover snapshot from RocksDB
        let (sm, snapshot, snapshot_idx) = Self::recover_from_disk(&engine);

        debug!(
            has_snapshot = snapshot.is_some(),
            snapshot_idx,
            "PersistentFlareStore initialized"
        );

        Self {
            engine,
            sm: RwLock::new(sm),
            snapshot_idx: RwLock::new(snapshot_idx),
            current_snapshot: RwLock::new(snapshot),
        }
    }

    /// Recover state from RocksDB on startup
    fn recover_from_disk(
        engine: &RocksEngine,
    ) -> (FlareStateMachine, Option<FlareSnapshot>, u64) {
        // Load snapshot index
        let snapshot_idx = engine
            .get_raft_state(KEY_SNAPSHOT_IDX)
            .ok()
            .flatten()
            .and_then(|data| {
                if data.len() == 8 {
                    let mut arr = [0u8; 8];
                    arr.copy_from_slice(&data);
                    Some(u64::from_be_bytes(arr))
                } else {
                    None
                }
            })
            .unwrap_or(0);

        // Load snapshot metadata
        let snapshot_meta: Option<SnapshotMeta<FlareNodeId, FlareNode>> = engine
            .get_raft_state(KEY_SNAPSHOT_META)
            .ok()
            .flatten()
            .and_then(|data| serde_json::from_slice(&data).ok());

        // Load snapshot data
        let snapshot_data: Option<Vec<u8>> = engine
            .get_raft_state(KEY_SNAPSHOT_DATA)
            .ok()
            .flatten();

        // If we have both meta and data, reconstruct the snapshot and state machine
        match (snapshot_meta, snapshot_data) {
            (Some(meta), Some(data)) => {
                // Try to deserialize state machine from snapshot (using bincode)
                match bincode::deserialize::<FlareStateMachine>(&data) {
                    Ok(sm) => {
                        debug!(
                            snapshot_id = %meta.snapshot_id,
                            "Recovered state machine from snapshot"
                        );
                        let snapshot = FlareSnapshot {
                            meta,
                            data,
                        };
                        (sm, Some(snapshot), snapshot_idx)
                    }
                    Err(e) => {
                        tracing::warn!(
                            error = %e,
                            "Failed to deserialize state machine from snapshot, starting fresh"
                        );
                        (FlareStateMachine::default(), None, 0)
                    }
                }
            }
            _ => {
                // No snapshot, start with empty state machine
                (FlareStateMachine::default(), None, snapshot_idx)
            }
        }
    }

    /// Persist snapshot to RocksDB
    fn persist_snapshot(&self, snapshot: &FlareSnapshot, idx: u64) -> Result<(), StorageError<FlareNodeId>> {
        // Persist snapshot metadata
        let meta_data = serde_json::to_vec(&snapshot.meta)
            .map_err(|e| StorageIOError::write(&e))?;
        self.engine
            .put_raft_state(KEY_SNAPSHOT_META, &meta_data)
            .map_err(|e| StorageIOError::write(&e))?;

        // Persist snapshot data
        self.engine
            .put_raft_state(KEY_SNAPSHOT_DATA, &snapshot.data)
            .map_err(|e| StorageIOError::write(&e))?;

        // Persist snapshot index
        self.engine
            .put_raft_state(KEY_SNAPSHOT_IDX, &idx.to_be_bytes())
            .map_err(|e| StorageIOError::write(&e))?;

        debug!(
            snapshot_id = %snapshot.meta.snapshot_id,
            idx,
            "Persisted snapshot to RocksDB"
        );

        Ok(())
    }

    /// Create a new PersistentFlareStore wrapped in Arc
    pub fn new_arc(engine: Arc<RocksEngine>) -> Arc<Self> {
        Arc::new(Self::new(engine))
    }

    /// Get the current state machine
    pub async fn get_state_machine(&self) -> FlareStateMachine {
        self.sm.read().await.clone()
    }

    /// Read a value from the state machine (eventual consistency)
    pub async fn read_kv(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64)> {
        let sm = self.sm.read().await;
        sm.kv_data.get(&(namespace_id, key.to_vec())).cloned()
    }

    /// Read a CAS value from the state machine (strong consistency)
    /// Returns (value, version, timestamp) if found
    pub async fn read_cas(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64, u64)> {
        let sm = self.sm.read().await;
        sm.cas_data.get(&(namespace_id, key.to_vec())).cloned()
    }

    /// Scan CAS data from the replicated state machine.
    pub async fn scan_cas(
        &self,
        namespace_id: u32,
        start_key: &[u8],
        end_key: &[u8],
        limit: usize,
    ) -> Vec<(Vec<u8>, Vec<u8>, u64, u64)> {
        let sm = self.sm.read().await;

        sm.cas_data
            .iter()
            .filter(|((ns_id, key), _)| {
                *ns_id == namespace_id
                    && key.as_slice() >= start_key
                    && (end_key.is_empty() || key.as_slice() < end_key)
            })
            .take(limit)
            .map(|((_ns_id, key), (value, version, ts))| {
                (key.clone(), value.clone(), *version, *ts)
            })
            .collect()
    }

    fn serialize_log_id(log_id: &FlareLogId) -> Vec<u8> {
        serde_json::to_vec(log_id).unwrap_or_default()
    }

    fn deserialize_log_id(data: &[u8]) -> Option<FlareLogId> {
        serde_json::from_slice(data).ok()
    }

    fn serialize_vote(vote: &Vote<FlareNodeId>) -> Vec<u8> {
        serde_json::to_vec(vote).unwrap_or_default()
    }

    fn deserialize_vote(data: &[u8]) -> Option<Vote<FlareNodeId>> {
        serde_json::from_slice(data).ok()
    }
}

impl RaftLogReader<FlareTypeConfig> for Arc<PersistentFlareStore> {
    async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + Debug + OptionalSend>(
        &mut self,
        range: RB,
    ) -> Result<Vec<FlareEntry>, StorageError<FlareNodeId>> {
        let start = match range.start_bound() {
            std::ops::Bound::Included(&n) => n,
            std::ops::Bound::Excluded(&n) => n + 1,
            std::ops::Bound::Unbounded => 0,
        };
        let end = match range.end_bound() {
            std::ops::Bound::Included(&n) => Some(n + 1),
            std::ops::Bound::Excluded(&n) => Some(n),
            std::ops::Bound::Unbounded => None,
        };

        let raw_entries = self
            .engine
            .get_raft_log_range(start, end)
            .map_err(|e| StorageIOError::read_logs(&e))?;

        let mut entries = Vec::with_capacity(raw_entries.len());
        for (_, data) in raw_entries {
            let entry: FlareEntry =
                serde_json::from_slice(&data).map_err(|e| StorageIOError::read_logs(&e))?;
            entries.push(entry);
        }

        Ok(entries)
    }
}

impl RaftSnapshotBuilder<FlareTypeConfig> for Arc<PersistentFlareStore> {
    async fn build_snapshot(
        &mut self,
    ) -> Result<Snapshot<FlareTypeConfig>, StorageError<FlareNodeId>> {
        let data;
        let last_applied_log;
        let last_membership;

        {
            let sm = self.sm.read().await;
            // Use bincode for state machine (handles tuple map keys)
            data = bincode::serialize(&*sm).map_err(|e| StorageIOError::read_state_machine(&e))?;
            last_applied_log = sm.last_applied_log;
            last_membership = sm.last_membership.clone();
        }

        let snapshot_idx = {
            let mut idx = self.snapshot_idx.write().await;
            *idx += 1;
            *idx
        };

        let snapshot_id = if let Some(last) = last_applied_log {
            format!("{}-{}-{}", last.leader_id, last.index, snapshot_idx)
        } else {
            format!("--{}", snapshot_idx)
        };

        let meta = SnapshotMeta {
            last_log_id: last_applied_log,
            last_membership,
            snapshot_id,
        };

        let snapshot = FlareSnapshot {
            meta: meta.clone(),
            data: data.clone(),
        };

        // Persist snapshot to RocksDB for recovery
        self.persist_snapshot(&snapshot, snapshot_idx)?;

        {
            let mut current = self.current_snapshot.write().await;
            *current = Some(snapshot);
        }

        Ok(Snapshot {
            meta,
            snapshot: Box::new(Cursor::new(data)),
        })
    }
}

impl RaftLogStorage<FlareTypeConfig> for Arc<PersistentFlareStore> {
    type LogReader = Self;

    async fn get_log_state(&mut self) -> Result<LogState<FlareTypeConfig>, StorageError<FlareNodeId>> {
        // Get last purged log ID from state
        let last_purged = self
            .engine
            .get_raft_state(KEY_LAST_PURGED)
            .map_err(|e| StorageIOError::read_logs(&e))?
            .and_then(|data| PersistentFlareStore::deserialize_log_id(&data));

        // Get last log entry from log
        let (_, last_idx) = self
            .engine
            .get_raft_log_bounds()
            .map_err(|e| StorageIOError::read_logs(&e))?;

        let last = if let Some(idx) = last_idx {
            let entry_data = self
                .engine
                .get_raft_log(idx)
                .map_err(|e| StorageIOError::read_logs(&e))?;

            if let Some(data) = entry_data {
                let entry: FlareEntry =
                    serde_json::from_slice(&data).map_err(|e| StorageIOError::read_logs(&e))?;
                Some(*entry.get_log_id())
            } else {
                last_purged
            }
        } else {
            last_purged
        };

        Ok(LogState {
            last_purged_log_id: last_purged,
            last_log_id: last,
        })
    }

    async fn save_vote(
        &mut self,
        vote: &Vote<FlareNodeId>,
    ) -> Result<(), StorageError<FlareNodeId>> {
        let data = PersistentFlareStore::serialize_vote(vote);
        self.engine
            .put_raft_state(KEY_VOTE, &data)
            .map_err(|e| StorageIOError::write_vote(&e))?;
        debug!(?vote, "Saved vote to persistent storage");
        Ok(())
    }

    async fn read_vote(&mut self) -> Result<Option<Vote<FlareNodeId>>, StorageError<FlareNodeId>> {
        let vote = self
            .engine
            .get_raft_state(KEY_VOTE)
            .map_err(|e| StorageIOError::read_vote(&e))?
            .and_then(|data| PersistentFlareStore::deserialize_vote(&data));
        Ok(vote)
    }

    async fn save_committed(
        &mut self,
        committed: Option<FlareLogId>,
    ) -> Result<(), StorageError<FlareNodeId>> {
        if let Some(log_id) = committed {
            let data = PersistentFlareStore::serialize_log_id(&log_id);
            self.engine
                .put_raft_state(KEY_COMMITTED, &data)
                .map_err(|e| StorageIOError::write(&e))?;
        }
        Ok(())
    }

    async fn read_committed(
        &mut self,
    ) -> Result<Option<FlareLogId>, StorageError<FlareNodeId>> {
        let committed = self
            .engine
            .get_raft_state(KEY_COMMITTED)
            .map_err(|e| StorageIOError::read(&e))?
            .and_then(|data| PersistentFlareStore::deserialize_log_id(&data));
        Ok(committed)
    }

    async fn append<I: IntoIterator<Item = FlareEntry> + Send>(
        &mut self,
        entries: I,
        callback: LogFlushed<FlareTypeConfig>,
    ) -> Result<(), StorageError<FlareNodeId>> {
        let entries: Vec<_> = entries.into_iter().collect();
        if entries.is_empty() {
            callback.log_io_completed(Ok(()));
            return Ok(());
        }

        for entry in &entries {
            let index = entry.log_id.index;
            let data = serde_json::to_vec(entry)
                .map_err(|e| StorageIOError::write_log_entry(*entry.get_log_id(), &e))?;
            self.engine
                .put_raft_log(index, &data)
                .map_err(|e| StorageIOError::write_log_entry(*entry.get_log_id(), &e))?;
            trace!(index, "Appended log entry to persistent storage");
        }

        callback.log_io_completed(Ok(()));
        Ok(())
    }

    async fn truncate(&mut self, log_id: FlareLogId) -> Result<(), StorageError<FlareNodeId>> {
        let (_, last_idx) = self
            .engine
            .get_raft_log_bounds()
            .map_err(|e| StorageIOError::write_logs(&e))?;

        if let Some(last) = last_idx {
            if log_id.index <= last {
                self.engine
                    .delete_raft_log_range(log_id.index, last)
                    .map_err(|e| StorageIOError::write_logs(&e))?;
            }
        }
        Ok(())
    }

    async fn purge(&mut self, log_id: FlareLogId) -> Result<(), StorageError<FlareNodeId>> {
        // Save last purged log ID
        let data = PersistentFlareStore::serialize_log_id(&log_id);
        self.engine
            .put_raft_state(KEY_LAST_PURGED, &data)
            .map_err(|e| StorageIOError::write_logs(&e))?;

        // Delete logs up to and including log_id.index
        let (first_idx, _) = self
            .engine
            .get_raft_log_bounds()
            .map_err(|e| StorageIOError::write_logs(&e))?;

        if let Some(first) = first_idx {
            if first <= log_id.index {
                self.engine
                    .delete_raft_log_range(first, log_id.index)
                    .map_err(|e| StorageIOError::write_logs(&e))?;
            }
        }

        Ok(())
    }

    async fn get_log_reader(&mut self) -> Self::LogReader {
        self.clone()
    }
}

impl RaftStateMachine<FlareTypeConfig> for Arc<PersistentFlareStore> {
    type SnapshotBuilder = Self;

    async fn applied_state(
        &mut self,
    ) -> Result<(Option<FlareLogId>, StoredMembership<FlareNodeId, FlareNode>), StorageError<FlareNodeId>> {
        let sm = self.sm.read().await;
        Ok((sm.last_applied_log, sm.last_membership.clone()))
    }

    async fn apply<I: IntoIterator<Item = FlareEntry> + Send>(
        &mut self,
        entries: I,
    ) -> Result<Vec<FlareResponse>, StorageError<FlareNodeId>> {
        let mut res = Vec::new();
        let mut sm = self.sm.write().await;

        for entry in entries {
            sm.last_applied_log = Some(entry.log_id);

            match &entry.payload {
                EntryPayload::Blank => {
                    res.push(FlareResponse::Ok);
                }
                EntryPayload::Normal(req) => {
                    let response = crate::storage::apply_request(&mut sm, req, entry.log_id.index);
                    res.push(response);
                }
                EntryPayload::Membership(mem) => {
                    sm.last_membership = StoredMembership::new(Some(entry.log_id), mem.clone());
                    res.push(FlareResponse::Ok);
                }
            }
        }

        Ok(res)
    }

    async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
        self.clone()
    }

    async fn begin_receiving_snapshot(
        &mut self,
    ) -> Result<Box<Cursor<Vec<u8>>>, StorageError<FlareNodeId>> {
        Ok(Box::new(Cursor::new(Vec::new())))
    }

    async fn install_snapshot(
        &mut self,
        meta: &SnapshotMeta<FlareNodeId, FlareNode>,
        snapshot: Box<Cursor<Vec<u8>>>,
    ) -> Result<(), StorageError<FlareNodeId>> {
        let new_snapshot = FlareSnapshot {
            meta: meta.clone(),
            data: snapshot.into_inner(),
        };

        // Update state machine (using bincode to match build_snapshot)
        {
            let new_sm: FlareStateMachine = bincode::deserialize(&new_snapshot.data)
                .map_err(|e| StorageIOError::read_snapshot(Some(new_snapshot.meta.signature()), &e))?;
            let mut sm = self.sm.write().await;
            *sm = new_sm;
        }

        // Persist snapshot to RocksDB for recovery
        let snapshot_idx = {
            let mut idx = self.snapshot_idx.write().await;
            *idx += 1;
            *idx
        };
        self.persist_snapshot(&new_snapshot, snapshot_idx)?;

        // Update current snapshot in memory
        let mut current = self.current_snapshot.write().await;
        *current = Some(new_snapshot);

        Ok(())
    }

    async fn get_current_snapshot(
        &mut self,
    ) -> Result<Option<Snapshot<FlareTypeConfig>>, StorageError<FlareNodeId>> {
        let current = self.current_snapshot.read().await;
        match &*current {
            Some(snapshot) => Ok(Some(Snapshot {
                meta: snapshot.meta.clone(),
                snapshot: Box::new(Cursor::new(snapshot.data.clone())),
            })),
            None => Ok(None),
        }
    }
}

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

    fn create_test_store() -> Arc<PersistentFlareStore> {
        let dir = tempdir().unwrap();
        let engine = Arc::new(RocksEngine::new(dir.path().to_str().unwrap()).unwrap());
        PersistentFlareStore::new_arc(engine)
    }

    #[tokio::test]
    async fn test_vote_persistence() {
        let mut store = create_test_store();

        let vote = Vote::new(5, 1);
        store.save_vote(&vote).await.unwrap();

        let loaded = store.read_vote().await.unwrap().unwrap();
        assert_eq!(loaded.leader_id().term, 5);
        assert_eq!(loaded.leader_id().node_id, 1);
    }

    #[tokio::test]
    async fn test_log_state_initial() {
        let mut store = create_test_store();

        let state = store.get_log_state().await.unwrap();
        assert!(state.last_log_id.is_none());
        assert!(state.last_purged_log_id.is_none());
    }

    #[tokio::test]
    async fn test_snapshot_persistence_and_recovery() {
        let dir = tempdir().unwrap();
        let path = dir.path().to_str().unwrap();

        // Create a store and populate state machine
        {
            let engine = Arc::new(RocksEngine::new(path).unwrap());
            let store = PersistentFlareStore::new_arc(engine);

            // Manually add data to state machine
            {
                let mut sm = store.sm.write().await;
                sm.kv_data
                    .insert((1, b"key1".to_vec()), (b"value1".to_vec(), 100));
                sm.cas_data
                    .insert((1, b"cas_key".to_vec()), (b"cas_value".to_vec(), 5, 200));
                sm.last_applied_log = Some(FlareLogId {
                    leader_id: openraft::LeaderId::new(1, 1),
                    index: 10,
                });
            }

            // Build and persist snapshot
            let mut store_mut = store.clone();
            let snapshot = store_mut.build_snapshot().await.unwrap();
            assert!(snapshot.meta.last_log_id.is_some());
        }

        // Reopen store and verify recovery
        {
            let engine = Arc::new(RocksEngine::new(path).unwrap());
            let store = PersistentFlareStore::new_arc(engine);

            // Verify state machine was recovered
            let kv = store.read_kv(1, b"key1").await;
            assert_eq!(kv, Some((b"value1".to_vec(), 100)));

            let cas = store.read_cas(1, b"cas_key").await;
            assert_eq!(cas, Some((b"cas_value".to_vec(), 5, 200)));

            // Verify snapshot is present
            let mut store_mut = store.clone();
            let current = store_mut.get_current_snapshot().await.unwrap();
            assert!(current.is_some());

            let sm = store.get_state_machine().await;
            assert!(sm.last_applied_log.is_some());
            assert_eq!(sm.last_applied_log.unwrap().index, 10);
        }
    }
}
