use crate::error::{Result, StorageError};
use crate::traits::KVStore;
use async_trait::async_trait;
use dashmap::DashMap;
use std::path::Path;
use std::sync::Arc;
use tracing::instrument;

fn fjall_err(e: fjall::Error) -> StorageError {
    StorageError::Backend(e.to_string())
}

/// Fjall (LSM-tree) implementation of the KVStore trait.
///
/// Used for write-heavy key prefixes: assertions (`H:`), votes (`V:`, `VC:`, `VW:`),
/// epochs (`E:`), supersession markers (`SUPERSEDED:`), and ingestion cursors (`__CURSOR__:`).
pub struct FjallStore {
    keyspace: fjall::Keyspace,
    partition: fjall::PartitionHandle,
    atomic_locks: Arc<DashMap<Vec<u8>, Arc<tokio::sync::Mutex<()>>>>,
    _temp_dir: Option<tempfile::TempDir>,
}

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

impl FjallStore {
    /// Open or create a Fjall database at the given path.
    #[instrument(skip_all)]
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let keyspace = fjall::Config::new(path.as_ref()).open().map_err(fjall_err)?;
        let partition = keyspace
            .open_partition("default", fjall::PartitionCreateOptions::default())
            .map_err(fjall_err)?;
        Ok(Self { keyspace, partition, atomic_locks: Arc::new(DashMap::new()), _temp_dir: None })
    }

    /// Open a temporary Fjall database for testing.
    ///
    /// The database will be automatically deleted when the returned store is dropped.
    pub fn open_temp() -> Result<Self> {
        let temp_dir = tempfile::tempdir().map_err(StorageError::Io)?;
        let keyspace = fjall::Config::new(temp_dir.path()).open().map_err(fjall_err)?;
        let partition = keyspace
            .open_partition("default", fjall::PartitionCreateOptions::default())
            .map_err(fjall_err)?;
        Ok(Self {
            keyspace,
            partition,
            atomic_locks: Arc::new(DashMap::new()),
            _temp_dir: Some(temp_dir),
        })
    }
}

#[async_trait]
impl KVStore for FjallStore {
    #[instrument(skip_all, fields(key_len = key.len()))]
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let result = self.partition.get(key).map_err(fjall_err)?;
        Ok(result.map(|slice| slice.to_vec()))
    }

    #[instrument(skip_all, fields(key_len = key.len(), value_len = value.len()))]
    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
        self.partition.insert(key, value).map_err(fjall_err)?;
        Ok(())
    }

    #[instrument(skip_all, fields(key_len = key.len()))]
    async fn delete(&self, key: &[u8]) -> Result<()> {
        self.partition.remove(key).map_err(fjall_err)?;
        Ok(())
    }

    #[instrument(skip_all, fields(prefix_len = prefix.len()))]
    async fn scan_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let mut results = Vec::new();
        for item in self.partition.prefix(prefix) {
            let (k, v) = item.map_err(fjall_err)?;
            results.push((k.to_vec(), v.to_vec()));
        }
        Ok(results)
    }

    #[instrument(skip_all)]
    async fn flush(&self) -> Result<()> {
        self.keyspace.persist(fjall::PersistMode::SyncAll).map_err(fjall_err)?;
        Ok(())
    }

    #[instrument(skip_all, fields(key_len = key.len(), delta))]
    async fn fetch_and_add_u64(&self, key: &[u8], delta: u64) -> Result<u64> {
        let lock = self
            .atomic_locks
            .entry(key.to_vec())
            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
            .clone();
        let _guard = lock.lock().await;

        let current = match self.partition.get(key).map_err(fjall_err)? {
            Some(bytes) => {
                let arr: [u8; 8] = bytes.as_ref().try_into().map_err(|_| {
                    StorageError::Serialization(format!(
                        "Corrupted u64 counter: expected 8 bytes, got {}",
                        bytes.len()
                    ))
                })?;
                u64::from_le_bytes(arr)
            }
            None => 0,
        };
        let new_val = current.saturating_add(delta);
        self.partition.insert(key, new_val.to_le_bytes()).map_err(fjall_err)?;
        Ok(new_val)
    }

    #[instrument(skip_all, fields(key_len = key.len()))]
    async fn compare_and_swap_f32<F>(&self, key: &[u8], update_fn: F) -> Result<f32>
    where
        F: Fn(f32) -> f32 + Send + Sync,
    {
        let lock = self
            .atomic_locks
            .entry(key.to_vec())
            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
            .clone();
        let _guard = lock.lock().await;

        let current = match self.partition.get(key).map_err(fjall_err)? {
            Some(bytes) => {
                let arr: [u8; 4] = bytes.as_ref().try_into().map_err(|_| {
                    StorageError::Serialization(format!(
                        "Corrupted f32 value: expected 4 bytes, got {}",
                        bytes.len()
                    ))
                })?;
                f32::from_le_bytes(arr)
            }
            None => 0.0,
        };
        let new_val = update_fn(current);
        self.partition.insert(key, new_val.to_le_bytes()).map_err(fjall_err)?;
        Ok(new_val)
    }
}

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

    #[tokio::test]
    async fn test_fjall_store_roundtrip() {
        let store = FjallStore::open_temp().expect("Failed to create temp DB");
        let key = b"test_key";
        let value = b"test_value";

        store.put(key, value).await.expect("Put failed");
        let retrieved = store.get(key).await.expect("Get failed");
        assert_eq!(retrieved, Some(value.to_vec()));

        store.delete(key).await.expect("Delete failed");
        let deleted = store.get(key).await.expect("Get failed");
        assert_eq!(deleted, None);
    }

    #[tokio::test]
    async fn test_fjall_scan_prefix() {
        let store = FjallStore::open_temp().expect("Failed to create temp DB");
        store.put(b"prefix:1", b"val1").await.unwrap();
        store.put(b"prefix:2", b"val2").await.unwrap();
        store.put(b"other:3", b"val3").await.unwrap();

        let results = store.scan_prefix(b"prefix:").await.unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0], (b"prefix:1".to_vec(), b"val1".to_vec()));
        assert_eq!(results[1], (b"prefix:2".to_vec(), b"val2".to_vec()));
    }

    #[tokio::test]
    async fn test_fjall_fetch_and_add() {
        let store = FjallStore::open_temp().expect("Failed to create temp DB");
        let key = b"counter";

        let val = store.fetch_and_add_u64(key, 5).await.unwrap();
        assert_eq!(val, 5);

        let val = store.fetch_and_add_u64(key, 3).await.unwrap();
        assert_eq!(val, 8);
    }

    #[tokio::test]
    async fn test_fjall_compare_and_swap_f32() {
        let store = FjallStore::open_temp().expect("Failed to create temp DB");
        let key = b"weight";

        let val = store.compare_and_swap_f32(key, |current| current + 1.5).await.unwrap();
        assert!((val - 1.5).abs() < f32::EPSILON);

        let val = store.compare_and_swap_f32(key, |current| current + 2.0).await.unwrap();
        assert!((val - 3.5).abs() < f32::EPSILON);
    }

    #[tokio::test]
    async fn test_fjall_flush() {
        let store = FjallStore::open_temp().expect("Failed to create temp DB");
        store.put(b"key", b"value").await.unwrap();
        store.flush().await.expect("Flush should succeed");
    }
}
