use std::collections::BTreeMap;
use std::fmt::Debug;
use std::io::{Cursor, Write};
use std::ops::RangeBounds;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::Mutex as StdMutex;

use openraft::storage::{
    LogFlushed, LogState, RaftLogReader, RaftLogStorage, RaftSnapshotBuilder, RaftStateMachine,
    Snapshot,
};
use openraft::{
    Entry, EntryPayload, LogId, OptionalSend, SnapshotMeta, StorageError, StorageIOError,
    StoredMembership, Vote,
};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;

use crate::failover::Decision;
use crate::rtype::{Node, NodeId, TypeConfig};

#[derive(Clone, Default)]
pub struct SharedDecision(Arc<StdMutex<Option<Decision>>>);

impl SharedDecision {
    pub fn current(&self) -> Option<Decision> {
        *self.0.lock().unwrap()
    }
    fn set(&self, decision: Option<Decision>) {
        *self.0.lock().unwrap() = decision;
    }
}

#[derive(Serialize, Deserialize, Default)]
struct LogStoreState {
    vote: Option<Vote<NodeId>>,
    last_purged: Option<LogId<NodeId>>,
    committed: Option<LogId<NodeId>>,
    log: BTreeMap<u64, Entry<TypeConfig>>,
}

#[derive(Clone)]
pub struct LogStore {
    path: PathBuf,
    state: Arc<RwLock<LogStoreState>>,
}

impl LogStore {
    async fn persist(&self) -> Result<(), StorageError<NodeId>> {
        let bytes = {
            let state = self.state.read().await;
            serde_json::to_vec(&*state).map_err(|e| StorageIOError::write_logs(&e))?
        };
        atomic_write(&self.path, &bytes).map_err(|e| StorageIOError::write_logs(&e))?;
        Ok(())
    }
}

impl RaftLogReader<TypeConfig> for LogStore {
    async fn try_get_log_entries<RB: RangeBounds<u64> + Clone + Debug + OptionalSend>(
        &mut self,
        range: RB,
    ) -> Result<Vec<Entry<TypeConfig>>, StorageError<NodeId>> {
        let state = self.state.read().await;
        Ok(state.log.range(range).map(|(_, entry)| entry.clone()).collect())
    }
}

impl RaftLogStorage<TypeConfig> for LogStore {
    type LogReader = Self;

    async fn get_log_state(&mut self) -> Result<LogState<TypeConfig>, StorageError<NodeId>> {
        let state = self.state.read().await;
        let last = state.log.iter().next_back().map(|(_, entry)| entry.log_id);
        let last_purged = state.last_purged;
        Ok(LogState {
            last_purged_log_id: last_purged,
            last_log_id: last.or(last_purged),
        })
    }

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

    async fn save_vote(&mut self, vote: &Vote<NodeId>) -> Result<(), StorageError<NodeId>> {
        {
            let mut state = self.state.write().await;
            state.vote = Some(*vote);
        }
        self.persist().await
    }

    async fn read_vote(&mut self) -> Result<Option<Vote<NodeId>>, StorageError<NodeId>> {
        Ok(self.state.read().await.vote)
    }

    async fn save_committed(
        &mut self,
        committed: Option<LogId<NodeId>>,
    ) -> Result<(), StorageError<NodeId>> {
        {
            let mut state = self.state.write().await;
            state.committed = committed;
        }
        self.persist().await
    }

    async fn read_committed(&mut self) -> Result<Option<LogId<NodeId>>, StorageError<NodeId>> {
        Ok(self.state.read().await.committed)
    }

    async fn append<I>(
        &mut self,
        entries: I,
        callback: LogFlushed<TypeConfig>,
    ) -> Result<(), StorageError<NodeId>>
    where
        I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
        I::IntoIter: OptionalSend,
    {
        {
            let mut state = self.state.write().await;
            for entry in entries {
                state.log.insert(entry.log_id.index, entry);
            }
        }
        self.persist().await?;
        callback.log_io_completed(Ok(()));
        Ok(())
    }

    async fn truncate(&mut self, log_id: LogId<NodeId>) -> Result<(), StorageError<NodeId>> {
        {
            let mut state = self.state.write().await;
            let keys: Vec<u64> = state.log.range(log_id.index..).map(|(k, _)| *k).collect();
            for key in keys {
                state.log.remove(&key);
            }
        }
        self.persist().await
    }

    async fn purge(&mut self, log_id: LogId<NodeId>) -> Result<(), StorageError<NodeId>> {
        {
            let mut state = self.state.write().await;
            state.last_purged = Some(log_id);
            let keys: Vec<u64> = state.log.range(..=log_id.index).map(|(k, _)| *k).collect();
            for key in keys {
                state.log.remove(&key);
            }
        }
        self.persist().await
    }
}

#[derive(Serialize, Deserialize, Default, Clone)]
struct SmState {
    last_applied: Option<LogId<NodeId>>,
    last_membership: StoredMembership<NodeId, Node>,
    decision: Option<Decision>,
}

struct StoredSnapshot {
    meta: SnapshotMeta<NodeId, Node>,
    data: Vec<u8>,
}

#[derive(Clone)]
pub struct StateMachine {
    path: PathBuf,
    sm: Arc<RwLock<SmState>>,
    snapshot: Arc<RwLock<Option<Arc<StoredSnapshot>>>>,
    snapshot_idx: Arc<StdMutex<u64>>,
    shared: SharedDecision,
}

impl StateMachine {
    async fn persist(&self) -> Result<(), StorageError<NodeId>> {
        let bytes = {
            let sm = self.sm.read().await;
            serde_json::to_vec(&*sm).map_err(|e| StorageIOError::write_state_machine(&e))?
        };
        atomic_write(&self.path, &bytes).map_err(|e| StorageIOError::write_state_machine(&e))?;
        Ok(())
    }
}

impl RaftSnapshotBuilder<TypeConfig> for StateMachine {
    async fn build_snapshot(&mut self) -> Result<Snapshot<TypeConfig>, StorageError<NodeId>> {
        let (data, last_applied, last_membership) = {
            let sm = self.sm.read().await;
            let data =
                serde_json::to_vec(&*sm).map_err(|e| StorageIOError::read_state_machine(&e))?;
            (data, sm.last_applied, sm.last_membership.clone())
        };
        let idx = {
            let mut guard = self.snapshot_idx.lock().unwrap();
            *guard += 1;
            *guard
        };
        let snapshot_id = match last_applied {
            Some(log_id) => format!("{}-{}-{}", log_id.leader_id, log_id.index, idx),
            None => format!("--{}", idx),
        };
        let meta = SnapshotMeta {
            last_log_id: last_applied,
            last_membership,
            snapshot_id,
        };
        {
            let mut current = self.snapshot.write().await;
            *current = Some(Arc::new(StoredSnapshot {
                meta: meta.clone(),
                data: data.clone(),
            }));
        }
        Ok(Snapshot {
            meta,
            snapshot: Box::new(Cursor::new(data)),
        })
    }
}

impl RaftStateMachine<TypeConfig> for StateMachine {
    type SnapshotBuilder = Self;

    async fn applied_state(
        &mut self,
    ) -> Result<(Option<LogId<NodeId>>, StoredMembership<NodeId, Node>), StorageError<NodeId>> {
        let sm = self.sm.read().await;
        Ok((sm.last_applied, sm.last_membership.clone()))
    }

    async fn apply<I>(&mut self, entries: I) -> Result<Vec<()>, StorageError<NodeId>>
    where
        I: IntoIterator<Item = Entry<TypeConfig>> + OptionalSend,
        I::IntoIter: OptionalSend,
    {
        let mut responses = Vec::new();
        let mut changed = false;
        let decision;
        {
            let mut sm = self.sm.write().await;
            for entry in entries {
                sm.last_applied = Some(entry.log_id);
                match entry.payload {
                    EntryPayload::Blank => {}
                    EntryPayload::Normal(value) => {
                        sm.decision = Some(value);
                        changed = true;
                    }
                    EntryPayload::Membership(membership) => {
                        sm.last_membership =
                            StoredMembership::new(Some(entry.log_id), membership);
                    }
                }
                responses.push(());
            }
            decision = sm.decision;
        }
        self.persist().await?;
        if changed {
            self.shared.set(decision);
        }
        Ok(responses)
    }

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

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

    async fn install_snapshot(
        &mut self,
        meta: &SnapshotMeta<NodeId, Node>,
        snapshot: Box<Cursor<Vec<u8>>>,
    ) -> Result<(), StorageError<NodeId>> {
        let data = snapshot.into_inner();
        let new_sm: SmState = serde_json::from_slice(&data)
            .map_err(|e| StorageIOError::read_snapshot(Some(meta.signature()), &e))?;
        let decision = new_sm.decision;
        {
            let mut sm = self.sm.write().await;
            *sm = new_sm;
        }
        {
            let mut current = self.snapshot.write().await;
            *current = Some(Arc::new(StoredSnapshot {
                meta: meta.clone(),
                data,
            }));
        }
        self.persist().await?;
        self.shared.set(decision);
        Ok(())
    }

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

pub fn open(dir: &Path, node_id: NodeId) -> (LogStore, StateMachine, SharedDecision) {
    let log_path = dir.join(format!("raft_log_{}.json", node_id));
    let sm_path = dir.join(format!("raft_sm_{}.json", node_id));
    let log_state: LogStoreState = load(&log_path);
    let sm_state: SmState = load(&sm_path);
    let shared = SharedDecision(Arc::new(StdMutex::new(sm_state.decision)));
    let log_store = LogStore {
        path: log_path,
        state: Arc::new(RwLock::new(log_state)),
    };
    let state_machine = StateMachine {
        path: sm_path,
        sm: Arc::new(RwLock::new(sm_state)),
        snapshot: Arc::new(RwLock::new(None)),
        snapshot_idx: Arc::new(StdMutex::new(0)),
        shared: shared.clone(),
    };
    (log_store, state_machine, shared)
}

fn load<T>(path: &Path) -> T
where
    T: Default + for<'de> Deserialize<'de>,
{
    match std::fs::read(path) {
        Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_else(|e| {
            panic!(
                "pg_replica: durable raft state at {} is present but unreadable ({}); refusing to start with empty state",
                path.display(),
                e
            )
        }),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => T::default(),
        Err(e) => panic!(
            "pg_replica: cannot read raft state at {} ({}); refusing to start",
            path.display(),
            e
        ),
    }
}

fn atomic_write(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
    let tmp = path.with_extension("tmp");
    {
        let mut file = std::fs::File::create(&tmp)?;
        file.write_all(bytes)?;
        file.sync_all()?;
    }
    std::fs::rename(&tmp, path)?;
    if let Some(dir) = path.parent().filter(|dir| !dir.as_os_str().is_empty()) {
        std::fs::File::open(dir)?.sync_all()?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use openraft::testing::{StoreBuilder, Suite};

    struct TestDir(PathBuf);

    impl Drop for TestDir {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    struct Builder;

    impl StoreBuilder<TypeConfig, LogStore, StateMachine, TestDir> for Builder {
        async fn build(&self) -> Result<(TestDir, LogStore, StateMachine), StorageError<NodeId>> {
            let nanos = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos();
            let dir = std::env::temp_dir().join(format!("pgr-store-test-{}-{}", std::process::id(), nanos));
            std::fs::create_dir_all(&dir).unwrap();
            let (log_store, state_machine, _shared) = open(&dir, 1);
            Ok((TestDir(dir), log_store, state_machine))
        }
    }

    #[test]
    fn storage_conformance_suite() -> Result<(), StorageError<NodeId>> {
        Suite::test_all(Builder)
    }
}