use redb::{Database, ReadableDatabase, TableDefinition};
use serde::{Serialize, de::DeserializeOwned};
use std::sync::{Arc, RwLock};

pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("store");

pub struct Store<T> {
    pub key: String,
    pub db: Arc<Database>,
    pub cache: RwLock<T>,
}

impl<T: DeserializeOwned + Default + Serialize + Clone + Send + 'static> Store<T> {
    pub fn new(key: &str, db: Arc<Database>) -> Self {
        let initial_data = Self::load_from_db(key, &db);
        Self {
            key: key.to_string(),
            db,
            cache: RwLock::new(initial_data),
        }
    }

    fn load_from_db(key: &str, db: &Database) -> T {
        let read_txn = db.begin_read().unwrap();
        if let Ok(table) = read_txn.open_table(STORE_TABLE) {
            if let Ok(Some(value)) = table.get(key) {
                if let Ok(parsed) = serde_json::from_slice::<T>(value.value()) {
                    return parsed;
                }
            }
        }
        T::default()
    }

    fn save_to_db(key: &str, db: &Database, data: &T) {
        let write_txn = db.begin_write().unwrap();
        {
            let mut table = write_txn.open_table(STORE_TABLE).unwrap();
            let json_data = serde_json::to_vec(data).unwrap();
            table.insert(key, json_data.as_slice()).unwrap();
        }
        write_txn.commit().unwrap();
    }

    pub fn read(&self) -> T {
        let lock = self.cache.read().unwrap();
        lock.clone()
    }

    pub fn modify<F: FnOnce(&mut T)>(&self, f: F) {
        let mut lock = self.cache.write().unwrap();
        f(&mut lock);
        let key = self.key.clone();
        let db = self.db.clone();
        let data = lock.clone();
        tokio::task::spawn_blocking(move || {
            Self::save_to_db(&key, &db, &data);
        });
    }
}

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

    #[derive(Serialize, serde::Deserialize, Clone, Default, PartialEq, Debug)]
    struct TestData {
        name: String,
        value: i32,
    }

    #[tokio::test]
    async fn test_store_read_write() {
        let temp_file = NamedTempFile::new().unwrap();
        let db = Database::create(temp_file.path()).unwrap();
        
        let write_txn = db.begin_write().unwrap();
        {
            write_txn.open_table(STORE_TABLE).unwrap();
        }
        write_txn.commit().unwrap();

        let db = Arc::new(db);
        let store = Store::<TestData>::new("test_key", db.clone());

        assert_eq!(store.read(), TestData::default());

        store.modify(|data| {
            data.name = "Hello".to_string();
            data.value = 42;
        });

        // Need to wait for spawn_blocking to finish
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        assert_eq!(store.read(), TestData { name: "Hello".to_string(), value: 42 });

        // Load again to verify persistence
        let store2 = Store::<TestData>::new("test_key", db.clone());
        assert_eq!(store2.read(), TestData { name: "Hello".to_string(), value: 42 });
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_store_concurrency() {
        let temp_file = NamedTempFile::new().unwrap();
        let db = Database::create(temp_file.path()).unwrap();
        
        let write_txn = db.begin_write().unwrap();
        {
            write_txn.open_table(STORE_TABLE).unwrap();
        }
        write_txn.commit().unwrap();

        let db = Arc::new(db);
        let store = Arc::new(Store::<TestData>::new("concurrent_key", db.clone()));
        
        let mut handles = vec![];
        for _ in 0..50 {
            let s = store.clone();
            handles.push(tokio::spawn(async move {
                s.modify(|data| {
                    data.value += 1;
                });
            }));
        }
        
        for h in handles {
            h.await.unwrap();
        }
        
        // Wait for all blocking writes to flush
        tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
        
        assert_eq!(store.read().value, 50);
    }
}