use std::collections::BTreeMap;
use std::sync::RwLock;

use super::WriteBatch;
use super::batch::BatchOp;
use super::engine::StorageEngine;
use super::error::StorageError;
use super::iterator::PrefixIterator;

/// In-memory storage backend backed by a `BTreeMap`.
///
/// Provides deterministic testing without disk I/O. All data lives
/// in memory and is lost when the backend is dropped.
///
/// Thread-safe via `RwLock` — concurrent reads are allowed,
/// writes are serialized.
pub struct InMemoryBackend {
    data: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
}

impl InMemoryBackend {
    #[must_use]
    pub const fn new() -> Self {
        Self {
            data: RwLock::new(BTreeMap::new()),
        }
    }
}

impl Default for InMemoryBackend {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for InMemoryBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("InMemoryBackend").finish_non_exhaustive()
    }
}

impl StorageEngine for InMemoryBackend {
    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
        let data = self.data.read().map_err(|_| StorageError::Closed)?;
        Ok(data.get(key).cloned())
    }

    fn put(&self, key: &[u8], value: &[u8]) -> Result<(), StorageError> {
        self.data
            .write()
            .map_err(|_| StorageError::Closed)?
            .insert(key.to_vec(), value.to_vec());
        Ok(())
    }

    fn delete(&self, key: &[u8]) -> Result<(), StorageError> {
        self.data
            .write()
            .map_err(|_| StorageError::Closed)?
            .remove(key);
        Ok(())
    }

    fn scan_prefix(&self, prefix: &[u8]) -> PrefixIterator<'_> {
        let Ok(data) = self.data.read() else {
            return Box::new(std::iter::once(Err(StorageError::Closed)));
        };

        // Collect matching entries into a Vec so we can drop the RwLockReadGuard
        // before returning the iterator — the guard cannot outlive this scope.
        let prefix_vec = prefix.to_vec();
        #[allow(clippy::needless_collect)]
        let entries: Vec<(Vec<u8>, Vec<u8>)> = data
            .range(prefix_vec.clone()..)
            .take_while(|(k, _)| k.starts_with(&prefix_vec))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        drop(data);

        Box::new(entries.into_iter().map(Ok))
    }

    fn write_batch(&self, batch: WriteBatch) -> Result<(), StorageError> {
        let mut data = self.data.write().map_err(|_| StorageError::Closed)?;
        for op in &batch.ops {
            match op {
                BatchOp::Put { key, value } => {
                    data.insert(key.clone(), value.clone());
                }
                BatchOp::Delete { key } => {
                    data.remove(key);
                }
            }
        }
        drop(data);
        Ok(())
    }

    fn flush(&self) -> Result<(), StorageError> {
        // No-op for in-memory backend.
        Ok(())
    }
}

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

    #[test]
    fn put_then_get() {
        let backend = InMemoryBackend::new();
        backend.put(b"key", b"value").unwrap();
        let val = backend.get(b"key").unwrap();
        assert_eq!(val.as_deref(), Some(b"value".as_slice()));
    }

    #[test]
    fn get_missing_returns_none() {
        let backend = InMemoryBackend::new();
        assert_eq!(backend.get(b"missing").unwrap(), None);
    }

    #[test]
    fn delete_removes_key() {
        let backend = InMemoryBackend::new();
        backend.put(b"key", b"value").unwrap();
        backend.delete(b"key").unwrap();
        assert_eq!(backend.get(b"key").unwrap(), None);
    }

    #[test]
    fn delete_nonexistent_is_ok() {
        let backend = InMemoryBackend::new();
        backend.delete(b"missing").unwrap();
    }

    #[test]
    fn overwrite_value() {
        let backend = InMemoryBackend::new();
        backend.put(b"key", b"v1").unwrap();
        backend.put(b"key", b"v2").unwrap();
        let val = backend.get(b"key").unwrap();
        assert_eq!(val.as_deref(), Some(b"v2".as_slice()));
    }

    #[test]
    fn scan_prefix_returns_matching() {
        let backend = InMemoryBackend::new();
        backend.put(b"abc1", b"v1").unwrap();
        backend.put(b"abc2", b"v2").unwrap();
        backend.put(b"abd1", b"v3").unwrap();
        backend.put(b"xyz", b"v4").unwrap();

        let results: Vec<_> = backend
            .scan_prefix(b"abc")
            .collect::<Result<Vec<_>, _>>()
            .unwrap();

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0, b"abc1");
        assert_eq!(results[1].0, b"abc2");
    }

    #[test]
    fn scan_prefix_empty_result() {
        let backend = InMemoryBackend::new();
        backend.put(b"abc", b"v1").unwrap();

        let results: Vec<_> = backend
            .scan_prefix(b"xyz")
            .collect::<Result<Vec<_>, _>>()
            .unwrap();

        assert!(results.is_empty());
    }

    #[test]
    fn write_batch_atomic() {
        let backend = InMemoryBackend::new();
        backend.put(b"existing", b"old").unwrap();

        let mut batch = WriteBatch::new();
        batch.put(b"key1".to_vec(), b"val1".to_vec());
        batch.put(b"key2".to_vec(), b"val2".to_vec());
        batch.delete(b"existing".to_vec());

        backend.write_batch(batch).unwrap();

        assert_eq!(
            backend.get(b"key1").unwrap().as_deref(),
            Some(b"val1".as_slice())
        );
        assert_eq!(
            backend.get(b"key2").unwrap().as_deref(),
            Some(b"val2".as_slice())
        );
        assert_eq!(backend.get(b"existing").unwrap(), None);
    }

    #[test]
    fn flush_is_noop() {
        let backend = InMemoryBackend::new();
        backend.flush().unwrap();
    }

    #[test]
    fn scan_prefix_lexicographic_order() {
        let backend = InMemoryBackend::new();
        // Insert in reverse order
        backend.put(b"pre_c", b"3").unwrap();
        backend.put(b"pre_a", b"1").unwrap();
        backend.put(b"pre_b", b"2").unwrap();

        let results: Vec<_> = backend
            .scan_prefix(b"pre_")
            .collect::<Result<Vec<_>, _>>()
            .unwrap();

        assert_eq!(results[0].0, b"pre_a");
        assert_eq!(results[1].0, b"pre_b");
        assert_eq!(results[2].0, b"pre_c");
    }

    mod proptests {
        use super::*;
        use proptest::prelude::*;

        proptest! {
            #[test]
            fn put_get_roundtrip(
                key in proptest::collection::vec(any::<u8>(), 1..100),
                value in proptest::collection::vec(any::<u8>(), 0..1000)
            ) {
                let backend = InMemoryBackend::new();
                backend.put(&key, &value).unwrap();
                let retrieved = backend.get(&key).unwrap().unwrap();
                prop_assert_eq!(retrieved, value);
            }

            #[test]
            fn delete_then_get_is_none(
                key in proptest::collection::vec(any::<u8>(), 1..100),
                value in proptest::collection::vec(any::<u8>(), 0..100)
            ) {
                let backend = InMemoryBackend::new();
                backend.put(&key, &value).unwrap();
                backend.delete(&key).unwrap();
                prop_assert_eq!(backend.get(&key).unwrap(), None);
            }

            #[test]
            fn scan_prefix_correctness(
                prefix in proptest::collection::vec(any::<u8>(), 1..10),
                suffixes in proptest::collection::vec(
                    proptest::collection::vec(any::<u8>(), 1..10),
                    1..20
                )
            ) {
                let backend = InMemoryBackend::new();

                // Insert keys with the prefix
                for suffix in &suffixes {
                    let mut key = prefix.clone();
                    key.extend_from_slice(suffix);
                    backend.put(&key, b"v").unwrap();
                }

                // Insert a key without the prefix
                let mut other = vec![0xFF];
                other.extend_from_slice(&prefix);
                backend.put(&other, b"other").unwrap();

                let results: Vec<_> = backend
                    .scan_prefix(&prefix)
                    .collect::<Result<Vec<_>, _>>()
                    .unwrap();

                // All results must start with the prefix
                for (k, _) in &results {
                    prop_assert!(k.starts_with(&prefix));
                }

                // Results must be sorted
                for window in results.windows(2) {
                    prop_assert!(window[0].0 <= window[1].0);
                }
            }
        }
    }
}
