use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::sync::RwLock;

use crate::keypath::timestamp_path;

/// Storage-layer errors.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StoreError {
    /// Save rejected because leader fence is missing/mismatched.
    NotLeader,
    /// Persisted timestamp payload does not match 8-byte compatibility format.
    InvalidTimestampEncoding {
        /// Actual byte length encountered while decoding.
        len: usize,
    },
    /// Attempted timestamp value is negative.
    NegativeTimestamp,
    /// Attempted write is not monotonic versus previously persisted timestamp.
    NonMonotonic {
        /// Previously persisted timestamp in nanoseconds.
        previous_nanos: i64,
        /// Attempted timestamp in nanoseconds.
        attempted_nanos: i64,
    },
    /// Internal storage backend failure.
    Internal(String),
}

impl Display for StoreError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotLeader => write!(f, "leadership has not been granted yet"),
            Self::InvalidTimestampEncoding { len } => {
                write!(f, "timestamp value must be 8 bytes, got {len}")
            }
            Self::NegativeTimestamp => write!(f, "timestamp nanos must be non-negative"),
            Self::NonMonotonic {
                previous_nanos,
                attempted_nanos,
            } => write!(
                f,
                "saving timestamp {attempted_nanos} is less than or equal to previous {previous_nanos}"
            ),
            Self::Internal(msg) => write!(f, "{msg}"),
        }
    }
}

impl std::error::Error for StoreError {}

/// Timestamp store abstraction used by allocator and cleaner paths.
pub trait TimestampStore: Send + Sync {
    /// Loads persisted timestamp nanos for one `(cluster_id, group_id)` key.
    fn load_timestamp(&self, cluster_id: u64, group_id: u32) -> Result<Option<i64>, StoreError>;

    /// Persists timestamp nanos if the provided leader fence still matches.
    fn save_timestamp_if_leader(
        &self,
        cluster_id: u64,
        group_id: u32,
        ts_nanos: i64,
        leader_key: &str,
        expected_leader_value: &str,
    ) -> Result<(), StoreError>;

    /// Deletes persisted timestamp for one `(cluster_id, group_id)` key.
    fn delete_timestamp(&self, cluster_id: u64, group_id: u32) -> Result<(), StoreError>;

    /// Records current leader value for fencing checks.
    fn set_leader_value(&self, leader_key: &str, leader_value: &str) -> Result<(), StoreError>;
}

/// In-memory etcd-like key-value store with transactional checks.
#[derive(Debug, Default)]
pub struct InMemoryEtcdStore {
    kv: RwLock<HashMap<String, Vec<u8>>>,
}

impl InMemoryEtcdStore {
    /// Creates an empty in-memory timestamp store.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets an arbitrary binary value; used by tests.
    pub fn put_raw(&self, key: impl Into<String>, value: Vec<u8>) -> Result<(), StoreError> {
        let mut kv = self
            .kv
            .write()
            .map_err(|_| StoreError::Internal("lock poisoned".to_string()))?;
        kv.insert(key.into(), value);
        Ok(())
    }

    fn load_raw(&self, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
        let kv = self
            .kv
            .read()
            .map_err(|_| StoreError::Internal("lock poisoned".to_string()))?;
        Ok(kv.get(key).cloned())
    }
}

impl TimestampStore for InMemoryEtcdStore {
    fn load_timestamp(&self, cluster_id: u64, group_id: u32) -> Result<Option<i64>, StoreError> {
        let key = timestamp_path(cluster_id, group_id);
        let raw = self.load_raw(&key)?;
        raw.map(|bytes| decode_timestamp_nanos(&bytes)).transpose()
    }

    fn save_timestamp_if_leader(
        &self,
        cluster_id: u64,
        group_id: u32,
        ts_nanos: i64,
        leader_key: &str,
        expected_leader_value: &str,
    ) -> Result<(), StoreError> {
        if expected_leader_value.is_empty() {
            return Err(StoreError::NotLeader);
        }
        if ts_nanos < 0 {
            return Err(StoreError::NegativeTimestamp);
        }

        let mut kv = self
            .kv
            .write()
            .map_err(|_| StoreError::Internal("lock poisoned".to_string()))?;

        let leader_value = kv
            .get(leader_key)
            .map(|v| String::from_utf8_lossy(v).to_string())
            .unwrap_or_default();
        if leader_value != expected_leader_value {
            return Err(StoreError::NotLeader);
        }

        let ts_key = timestamp_path(cluster_id, group_id);
        if let Some(prev_raw) = kv.get(&ts_key) {
            let previous_nanos = decode_timestamp_nanos(prev_raw)?;
            if ts_nanos <= previous_nanos {
                return Err(StoreError::NonMonotonic {
                    previous_nanos,
                    attempted_nanos: ts_nanos,
                });
            }
        }

        kv.insert(ts_key, encode_timestamp_nanos(ts_nanos).to_vec());
        Ok(())
    }

    fn delete_timestamp(&self, cluster_id: u64, group_id: u32) -> Result<(), StoreError> {
        let key = timestamp_path(cluster_id, group_id);
        let mut kv = self
            .kv
            .write()
            .map_err(|_| StoreError::Internal("lock poisoned".to_string()))?;
        kv.remove(&key);
        Ok(())
    }

    fn set_leader_value(&self, leader_key: &str, leader_value: &str) -> Result<(), StoreError> {
        let mut kv = self
            .kv
            .write()
            .map_err(|_| StoreError::Internal("lock poisoned".to_string()))?;
        kv.insert(leader_key.to_string(), leader_value.as_bytes().to_vec());
        Ok(())
    }
}

/// Encodes Unix-nanoseconds timestamp into fixed-width big-endian bytes.
#[must_use]
pub fn encode_timestamp_nanos(nanos: i64) -> [u8; 8] {
    (nanos as u64).to_be_bytes()
}

/// Decodes fixed-width big-endian bytes into Unix-nanoseconds timestamp.
pub fn decode_timestamp_nanos(raw: &[u8]) -> Result<i64, StoreError> {
    if raw.len() != 8 {
        return Err(StoreError::InvalidTimestampEncoding { len: raw.len() });
    }
    let mut bytes = [0_u8; 8];
    bytes.copy_from_slice(raw);
    Ok(u64::from_be_bytes(bytes) as i64)
}

#[cfg(test)]
mod tests {
    use crate::keypath::{election_path, timestamp_path};

    use super::{
        InMemoryEtcdStore, StoreError, TimestampStore, decode_timestamp_nanos,
        encode_timestamp_nanos,
    };

    #[test]
    fn timestamp_encoding_matches_contract_vectors() {
        assert_eq!(encode_timestamp_nanos(1), 1_u64.to_be_bytes());
        assert_eq!(
            encode_timestamp_nanos(1_000_000),
            0x0000_0000_000f_4240_u64.to_be_bytes()
        );
    }

    #[test]
    fn decode_rejects_non_8_bytes() {
        let err = decode_timestamp_nanos(&[1, 2, 3]).expect_err("must fail");
        assert!(matches!(
            err,
            StoreError::InvalidTimestampEncoding { len: 3 }
        ));
    }

    #[test]
    fn fenced_save_requires_leader_and_monotonicity() {
        let store = InMemoryEtcdStore::new();
        let cluster_id = 42;
        let group_id = 7;
        let leader_key = election_path(cluster_id, group_id);

        let err = store
            .save_timestamp_if_leader(cluster_id, group_id, 10, &leader_key, "")
            .expect_err("empty leader must fail");
        assert!(matches!(err, StoreError::NotLeader));

        store
            .set_leader_value(&leader_key, "node-a")
            .expect("set leader");
        store
            .save_timestamp_if_leader(cluster_id, group_id, 100, &leader_key, "node-a")
            .expect("first save");

        let err = store
            .save_timestamp_if_leader(cluster_id, group_id, 100, &leader_key, "node-a")
            .expect_err("non-monotonic save must fail");
        assert!(matches!(err, StoreError::NonMonotonic { .. }));

        let key = timestamp_path(cluster_id, group_id);
        let loaded = store.load_raw(&key).expect("read raw").expect("present");
        assert_eq!(loaded, encode_timestamp_nanos(100).to_vec());
    }
}