use crate::{CasResult, StorageEngine, StorageError};
use async_trait::async_trait;
use rocksdb::{ColumnFamilyDescriptor, Options, DB};
use std::sync::Arc;

pub struct RocksEngine {
    db: Arc<DB>,
}

impl RocksEngine {
    pub fn new(path: &str) -> Result<Self, StorageError> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cf_opts = Options::default();
        let cfs = vec![
            ColumnFamilyDescriptor::new("default", cf_opts.clone()),
            ColumnFamilyDescriptor::new("cas", cf_opts.clone()),
            ColumnFamilyDescriptor::new("raft_log", cf_opts.clone()),
            ColumnFamilyDescriptor::new("raft_state", cf_opts),
        ];

        let db = DB::open_cf_descriptors(&opts, path, cfs)?;
        Ok(Self { db: Arc::new(db) })
    }

    pub fn db(&self) -> Arc<DB> {
        Arc::clone(&self.db)
    }

    /// Approximate bytes between [start, end). Used for split threshold decisions.
    pub fn approximate_range_size(&self, start: &[u8], end: &[u8]) -> Result<u64, StorageError> {
        let cf = self
            .db
            .cf_handle("default")
            .ok_or_else(|| StorageError::Serialization("CF default not found".to_string()))?;
        let mut iter = self.db.iterator_cf(
            &cf,
            rocksdb::IteratorMode::From(start, rocksdb::Direction::Forward),
        );
        let mut total: u64 = 0;
        while let Some(Ok((k, v))) = iter.next() {
            if !end.is_empty() && k.as_ref() >= end {
                break;
            }
            total = total.saturating_add(k.len() as u64 + v.len() as u64);
        }
        Ok(total)
    }

    fn get_version_from_val(val: &[u8]) -> Result<u64, StorageError> {
        if val.len() < 8 {
            return Err(StorageError::Serialization(
                "Value too short for version".to_string(),
            ));
        }
        let mut bytes = [0u8; 8];
        bytes.copy_from_slice(&val[0..8]);
        Ok(u64::from_be_bytes(bytes))
    }

    fn split_cas_value(val: &[u8]) -> Result<(u64, Vec<u8>), StorageError> {
        let version = Self::get_version_from_val(val)?;
        Ok((version, val[8..].to_vec()))
    }

    // ==================== Raft Log Storage ====================

    /// Get a Raft log entry by index
    pub fn get_raft_log(&self, index: u64) -> Result<Option<Vec<u8>>, StorageError> {
        let cf = self
            .db
            .cf_handle("raft_log")
            .ok_or_else(|| StorageError::Serialization("CF raft_log not found".to_string()))?;
        let key = index.to_be_bytes();
        self.db.get_cf(&cf, key).map_err(StorageError::RocksDb)
    }

    /// Put a Raft log entry
    pub fn put_raft_log(&self, index: u64, entry: &[u8]) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("raft_log")
            .ok_or_else(|| StorageError::Serialization("CF raft_log not found".to_string()))?;
        let key = index.to_be_bytes();
        self.db
            .put_cf(&cf, key, entry)
            .map_err(StorageError::RocksDb)
    }

    /// Delete Raft log entries from start_index to end_index (inclusive)
    pub fn delete_raft_log_range(
        &self,
        start_index: u64,
        end_index: u64,
    ) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("raft_log")
            .ok_or_else(|| StorageError::Serialization("CF raft_log not found".to_string()))?;
        for index in start_index..=end_index {
            let key = index.to_be_bytes();
            self.db.delete_cf(&cf, key).map_err(StorageError::RocksDb)?;
        }
        Ok(())
    }

    /// Get Raft log entries in range [start, end)
    pub fn get_raft_log_range(
        &self,
        start: u64,
        end: Option<u64>,
    ) -> Result<Vec<(u64, Vec<u8>)>, StorageError> {
        let cf = self
            .db
            .cf_handle("raft_log")
            .ok_or_else(|| StorageError::Serialization("CF raft_log not found".to_string()))?;

        let start_key = start.to_be_bytes();
        let iter = self.db.iterator_cf(
            &cf,
            rocksdb::IteratorMode::From(&start_key, rocksdb::Direction::Forward),
        );

        let mut entries = Vec::new();
        for item in iter {
            let (key_bytes, value) = item.map_err(StorageError::RocksDb)?;
            if key_bytes.len() != 8 {
                continue;
            }
            let mut key_arr = [0u8; 8];
            key_arr.copy_from_slice(&key_bytes);
            let index = u64::from_be_bytes(key_arr);

            if let Some(end_idx) = end {
                if index >= end_idx {
                    break;
                }
            }
            entries.push((index, value.to_vec()));
        }
        Ok(entries)
    }

    /// Get first and last log index
    pub fn get_raft_log_bounds(&self) -> Result<(Option<u64>, Option<u64>), StorageError> {
        let cf = self
            .db
            .cf_handle("raft_log")
            .ok_or_else(|| StorageError::Serialization("CF raft_log not found".to_string()))?;

        // Get first
        let first = {
            let mut iter = self
                .db
                .iterator_cf(&cf, rocksdb::IteratorMode::Start);
            iter.next()
                .and_then(|r| r.ok())
                .and_then(|(k, _)| {
                    if k.len() == 8 {
                        let mut arr = [0u8; 8];
                        arr.copy_from_slice(&k);
                        Some(u64::from_be_bytes(arr))
                    } else {
                        None
                    }
                })
        };

        // Get last
        let last = {
            let mut iter = self.db.iterator_cf(&cf, rocksdb::IteratorMode::End);
            iter.next()
                .and_then(|r| r.ok())
                .and_then(|(k, _)| {
                    if k.len() == 8 {
                        let mut arr = [0u8; 8];
                        arr.copy_from_slice(&k);
                        Some(u64::from_be_bytes(arr))
                    } else {
                        None
                    }
                })
        };

        Ok((first, last))
    }

    // ==================== Raft State Storage ====================

    /// Get Raft state value by key
    pub fn get_raft_state(&self, key: &str) -> Result<Option<Vec<u8>>, StorageError> {
        let cf = self
            .db
            .cf_handle("raft_state")
            .ok_or_else(|| StorageError::Serialization("CF raft_state not found".to_string()))?;
        self.db
            .get_cf(&cf, key.as_bytes())
            .map_err(StorageError::RocksDb)
    }

    /// Put Raft state value
    pub fn put_raft_state(&self, key: &str, value: &[u8]) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("raft_state")
            .ok_or_else(|| StorageError::Serialization("CF raft_state not found".to_string()))?;
        self.db
            .put_cf(&cf, key.as_bytes(), value)
            .map_err(StorageError::RocksDb)
    }

    /// Delete Raft state value
    pub fn delete_raft_state(&self, key: &str) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("raft_state")
            .ok_or_else(|| StorageError::Serialization("CF raft_state not found".to_string()))?;
        self.db
            .delete_cf(&cf, key.as_bytes())
            .map_err(StorageError::RocksDb)
    }
}

#[async_trait]
impl StorageEngine for RocksEngine {
    async fn put_raw(&self, key: &[u8], value: &[u8]) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("default")
            .ok_or_else(|| StorageError::Serialization("CF default not found".to_string()))?;

        self.db
            .put_cf(&cf, key, value)
            .map_err(StorageError::RocksDb)
    }

    async fn get_raw(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
        let cf = self
            .db
            .cf_handle("default")
            .ok_or_else(|| StorageError::Serialization("CF default not found".to_string()))?;
        self.db.get_cf(&cf, key).map_err(StorageError::RocksDb)
    }

    async fn delete_raw(&self, key: &[u8]) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("default")
            .ok_or_else(|| StorageError::Serialization("CF default not found".to_string()))?;
        self.db
            .delete_cf(&cf, key)
            .map_err(StorageError::RocksDb)
    }

    async fn get_cas(&self, key: &[u8]) -> Result<Option<(u64, Vec<u8>)>, StorageError> {
        let cf = self
            .db
            .cf_handle("cas")
            .ok_or_else(|| StorageError::Serialization("CF cas not found".to_string()))?;

        let raw = self.db.get_cf(&cf, key)?;
        match raw {
            Some(bytes) => Ok(Some(Self::split_cas_value(&bytes)?)),
            None => Ok(None),
        }
    }

    async fn compare_and_swap(
        &self,
        key: &[u8],
        expected_version: u64,
        value: &[u8],
    ) -> Result<CasResult, StorageError> {
        let cf = self
            .db
            .cf_handle("cas")
            .ok_or_else(|| StorageError::Serialization("CF cas not found".to_string()))?;

        // Note: In a real distributed system with concurrent local threads, we'd need a mutex or RocksDB TransactionDB.
        // For this MVP (single-node, or single-threaded access per key assumption for now), simple get/put logic is implemented.
        // Since we are inside async context, we should be careful about blocking. RocksDB is blocking.
        // For true atomic CAS in RocksDB without TransactionDB, we'd use MergeOperator or a custom C++ extension.
        // Here we implement a "pessimistic" check-then-act which is safe ONLY if we strictly serialize access to this key.
        // FlareDB architecture assumes Single-Leader (Raft) or serial execution for a given key range, so this is "acceptable" for Phase 1 MVP.

        // 1. Get current
        let current_val_opt = self.db.get_cf(&cf, key)?;
        let current_version = match &current_val_opt {
            Some(v) => Self::get_version_from_val(v)?,
            None => 0,
        };

        if current_version != expected_version {
            return Ok(CasResult {
                success: false,
                current_version,
                new_version: 0,
            });
        }

        let new_version = current_version + 1;
        let mut new_val_bytes = Vec::with_capacity(8 + value.len());
        new_val_bytes.extend_from_slice(&new_version.to_be_bytes());
        new_val_bytes.extend_from_slice(value);

        self.db.put_cf(&cf, key, &new_val_bytes)?;

        Ok(CasResult {
            success: true,
            current_version,
            new_version,
        })
    }

    async fn delete_cas(&self, key: &[u8]) -> Result<(), StorageError> {
        let cf = self
            .db
            .cf_handle("cas")
            .ok_or_else(|| StorageError::Serialization("CF cas not found".to_string()))?;
        self.db
            .delete_cf(&cf, key)
            .map_err(StorageError::RocksDb)
    }

    async fn scan_raw(
        &self,
        start: &[u8],
        end: &[u8],
        limit: usize,
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
        let cf = self
            .db
            .cf_handle("default")
            .ok_or_else(|| StorageError::Serialization("CF default not found".to_string()))?;

        let iter = self.db.iterator_cf(
            &cf,
            rocksdb::IteratorMode::From(start, rocksdb::Direction::Forward),
        );

        let mut entries = Vec::new();
        for item in iter {
            if entries.len() >= limit {
                break;
            }
            let (key, value) = item.map_err(StorageError::RocksDb)?;
            // Stop if we've passed the end key
            if !end.is_empty() && key.as_ref() >= end {
                break;
            }
            entries.push((key.to_vec(), value.to_vec()));
        }
        Ok(entries)
    }

    async fn scan_cas(
        &self,
        start: &[u8],
        end: &[u8],
        limit: usize,
    ) -> Result<Vec<(Vec<u8>, Vec<u8>, u64)>, StorageError> {
        let cf = self
            .db
            .cf_handle("cas")
            .ok_or_else(|| StorageError::Serialization("CF cas not found".to_string()))?;

        let iter = self.db.iterator_cf(
            &cf,
            rocksdb::IteratorMode::From(start, rocksdb::Direction::Forward),
        );

        let mut entries = Vec::new();
        for item in iter {
            if entries.len() >= limit {
                break;
            }
            let (key, value) = item.map_err(StorageError::RocksDb)?;
            // Stop if we've passed the end key
            if !end.is_empty() && key.as_ref() >= end {
                break;
            }
            let (version, data) = Self::split_cas_value(&value)?;
            entries.push((key.to_vec(), data, version));
        }
        Ok(entries)
    }
}
