use anyhow::{Result, Context};
use async_trait::async_trait;
use uuid::Uuid;
use chrono::Utc;
use memrec_common::{Memory, MemoryType};
use super::traits::MemoryStorage;
use super::rocksdb::RocksDBStore;

pub struct MemoryStore {
    rocksdb: RocksDBStore,
}

impl MemoryStore {
    pub fn new(rocksdb: RocksDBStore) -> Self {
        Self { rocksdb }
    }
    
    fn memory_key(id: &Uuid) -> Vec<u8> {
        id.to_string().into_bytes()
    }
    
    fn tag_key(tag: &str, id: &Uuid) -> Vec<u8> {
        format!("{}:{}", tag, id).into_bytes()
    }
    
    fn serialize_memory(memory: &Memory) -> Result<Vec<u8>> {
        serde_json::to_vec(memory).context("Failed to serialize memory")
    }
    
    fn deserialize_memory(data: &[u8]) -> Result<Memory> {
        serde_json::from_slice(data).context("Failed to deserialize memory")
    }
    
    fn collect_memories_from_iter(
        mut iter: rocksdb::DBRawIterator<'_>,
        limit: usize,
        filter_deleted: bool,
    ) -> Vec<Memory> {
        let mut memories = Vec::new();
        let mut count = 0;
        
        iter.seek_to_first();
        while iter.valid() && count < limit {
            if let Some(value) = iter.value() {
                if let Ok(memory) = Self::deserialize_memory(value) {
                    if !filter_deleted || !memory.is_deleted {
                        memories.push(memory);
                        count += 1;
                    }
                }
            }
            iter.next();
        }
        
        memories
    }
}

#[async_trait]
impl MemoryStorage for MemoryStore {
    async fn save(&self, memory: &Memory) -> Result<()> {
        let id_key = Self::memory_key(&memory.id);
        let data = Self::serialize_memory(memory)?;
        
        let cf_memories = self.rocksdb.cf_memories()?;
        self.rocksdb.put_cf(cf_memories, &id_key, &data)?;
        
        for tag in &memory.tags {
            let tag_key = Self::tag_key(tag, &memory.id);
            let cf_by_tag = self.rocksdb.cf_by_tag()?;
            self.rocksdb.put_cf(cf_by_tag, &tag_key, &id_key)?;
        }
        
        let importance_key = Self::memory_key(&memory.id);
        let importance_data = memory.importance.to_string().into_bytes();
        let cf_importance = self.rocksdb.cf_importance()?;
        self.rocksdb.put_cf(cf_importance, &importance_key, &importance_data)?;
        
        Ok(())
    }
    
    async fn get(&self, id: &Uuid) -> Result<Option<Memory>> {
        let id_key = Self::memory_key(id);
        let cf_memories = self.rocksdb.cf_memories()?;
        
        match self.rocksdb.get_cf(cf_memories, &id_key)? {
            Some(bytes) => {
                let memory = Self::deserialize_memory(&bytes)?;
                Ok(Some(memory))
            }
            None => Ok(None)
        }
    }
    
    async fn update(&self, memory: &Memory) -> Result<()> {
        self.save(memory).await
    }
    
    async fn delete(&self, id: &Uuid) -> Result<bool> {
        let memory = self.get(id).await?;
        
        match memory {
            Some(mem) => {
                if mem.is_deleted {
                    let id_key = Self::memory_key(id);
                    let cf_memories = self.rocksdb.cf_memories()?;
                    self.rocksdb.delete_cf(cf_memories, &id_key)?;
                    
                    let cf_deleted = self.rocksdb.cf_deleted()?;
                    self.rocksdb.delete_cf(cf_deleted, &id_key)?;
                    
                    for tag in &mem.tags {
                        let tag_key = Self::tag_key(tag, id);
                        let cf_by_tag = self.rocksdb.cf_by_tag()?;
                        self.rocksdb.delete_cf(cf_by_tag, &tag_key)?;
                    }
                    
                    Ok(true)
                } else {
                    let mut mem_deleted = mem;
                    mem_deleted.is_deleted = true;
                    mem_deleted.deleted_at = Some(Utc::now());
                    self.update(&mem_deleted).await?;
                    
                    let id_key = Self::memory_key(id);
                    let cf_deleted = self.rocksdb.cf_deleted()?;
                    self.rocksdb.put_cf(cf_deleted, &id_key, &id_key)?;
                    
                    Ok(false)
                }
            }
            None => Ok(false)
        }
    }
    
    async fn list(&self, limit: usize) -> Result<Vec<Memory>> {
        let cf_memories = self.rocksdb.cf_memories()?;
        let iter = self.rocksdb.iter_cf(cf_memories);
        
        Ok(Self::collect_memories_from_iter(iter, limit, true))
    }
    
    async fn list_by_type(&self, memory_type: MemoryType, limit: usize) -> Result<Vec<Memory>> {
        let all = self.list(limit * 10).await?;
        Ok(all.into_iter()
            .filter(|m| m.memory_type == memory_type)
            .take(limit)
            .collect())
    }
    
    async fn list_by_tag(&self, tag: &str, limit: usize) -> Result<Vec<Memory>> {
        let cf_by_tag = self.rocksdb.cf_by_tag()?;
        let mut iter = self.rocksdb.iter_cf(cf_by_tag);
        
        let prefix = format!("{}:", tag).into_bytes();
        let mut memories = Vec::new();
        let mut count = 0;
        
        iter.seek_to_first();
        while iter.valid() && count < limit {
            if let Some(key) = iter.key() {
                if key.starts_with(&prefix) {
                    if let Some(value) = iter.value() {
                        let id_str = String::from_utf8_lossy(value);
                        if let Ok(id) = Uuid::parse_str(&id_str) {
                            if let Some(memory) = self.get(&id).await? {
                                if !memory.is_deleted {
                                    memories.push(memory);
                                    count += 1;
                                }
                            }
                        }
                    }
                }
            }
            iter.next();
        }
        
        Ok(memories)
    }
    
    async fn list_by_importance(&self, min: f32, max: f32) -> Result<Vec<Memory>> {
        let cf_importance = self.rocksdb.cf_importance()?;
        let mut iter = self.rocksdb.iter_cf(cf_importance);
        
        let mut memories = Vec::new();
        
        iter.seek_to_first();
        while iter.valid() {
            if let Some(key) = iter.key() {
                if let Some(value) = iter.value() {
                    let importance_str = String::from_utf8_lossy(value);
                    if let Ok(importance) = importance_str.parse::<f32>() {
                        if importance >= min && importance <= max {
                            let id_str = String::from_utf8_lossy(key);
                            if let Ok(id) = Uuid::parse_str(&id_str) {
                                if let Some(memory) = self.get(&id).await? {
                                    if !memory.is_deleted {
                                        memories.push(memory);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            iter.next();
        }
        
        Ok(memories)
    }
    
    async fn list_deleted(&self) -> Result<Vec<Memory>> {
        let cf_deleted = self.rocksdb.cf_deleted()?;
        let mut iter = self.rocksdb.iter_cf(cf_deleted);
        
        let mut memories = Vec::new();
        
        iter.seek_to_first();
        while iter.valid() {
            if let Some(key) = iter.key() {
                let id_str = String::from_utf8_lossy(key);
                if let Ok(id) = Uuid::parse_str(&id_str) {
                    if let Some(memory) = self.get(&id).await? {
                        if memory.is_deleted {
                            memories.push(memory);
                        }
                    }
                }
            }
            iter.next();
        }
        
        Ok(memories)
    }
    
    async fn count(&self) -> Result<usize> {
        let memories = self.list(10000).await?;
        Ok(memories.len())
    }
    
    async fn count_deleted(&self) -> Result<usize> {
        let deleted = self.list_deleted().await?;
        Ok(deleted.len())
    }
    
    async fn list_by_project(&self, project_id: &Uuid) -> Result<Vec<Memory>> {
        let all = self.list(10000).await?;
        Ok(all.into_iter()
            .filter(|m| m.project_id == Some(*project_id))
            .collect())
    }
    
    async fn get_chunks_by_group(&self, chunk_group_id: &Uuid) -> Result<Vec<Memory>> {
        let all = self.list(10000).await?;
        Ok(all.into_iter()
            .filter(|m| m.chunk_group_id == Some(*chunk_group_id))
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    
    #[tokio::test]
    async fn test_memory_save_and_get() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(rocksdb);
        
        let memory = Memory::new("test content".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["tag1".to_string()]);
        
        store.save(&memory).await.unwrap();
        
        let retrieved = store.get(&memory.id).await.unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().content, "test content");
    }
    
    #[tokio::test]
    async fn test_memory_list() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(rocksdb);
        
        for i in 0..5 {
            let memory = Memory::new(format!("content {}", i), MemoryType::Conversation);
            store.save(&memory).await.unwrap();
        }
        
        let memories = store.list(10).await.unwrap();
        assert_eq!(memories.len(), 5);
    }
    
    #[tokio::test]
    async fn test_memory_list_by_tag() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(rocksdb);
        
        let m1 = Memory::new("content 1".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["rust".to_string()]);
        let m2 = Memory::new("content 2".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["python".to_string()]);
        
        store.save(&m1).await.unwrap();
        store.save(&m2).await.unwrap();
        
        let rust_memories = store.list_by_tag("rust", 10).await.unwrap();
        assert_eq!(rust_memories.len(), 1);
    }
    
    #[tokio::test]
    async fn test_memory_soft_delete() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(rocksdb);
        
        let memory = Memory::new("test".to_string(), MemoryType::Conversation);
        store.save(&memory).await.unwrap();
        
        let deleted = store.delete(&memory.id).await.unwrap();
        assert!(!deleted);
        
        let retrieved = store.get(&memory.id).await.unwrap();
        assert!(retrieved.is_some());
        assert!(retrieved.unwrap().is_deleted);
    }
}