use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use chrono::{DateTime, Utc};
use serde::{Serialize, Deserialize};
use crate::{Envelope, UCXLUri, Result, UCXLError, EnvelopeHistory, VersionInfo};

#[async_trait::async_trait]
pub trait EnvelopeStore: Send + Sync {
    async fn store(&self, envelope: &Envelope) -> Result<()>;
    async fn retrieve(&self, envelope_id: &str) -> Result<Option<Envelope>>;
    async fn retrieve_by_uri(&self, uri: &UCXLUri) -> Result<Option<Envelope>>;
    async fn retrieve_version(&self, envelope_id: &str, version: &str) -> Result<Option<Envelope>>;
    async fn list_versions(&self, envelope_id: &str) -> Result<Vec<VersionInfo>>;
    async fn get_history(&self, doc_id: &str) -> Result<EnvelopeHistory>;
    async fn delete(&self, envelope_id: &str, soft_delete: bool) -> Result<()>;
    async fn search(&self, query: &SearchQuery) -> Result<Vec<Envelope>>;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchQuery {
    pub text: Option<String>,
    pub tags: Vec<String>,
    pub author: Option<String>,
    pub content_type: Option<String>,
    pub date_range: Option<DateRange>,
    pub limit: Option<usize>,
    pub offset: Option<usize>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DateRange {
    pub from: DateTime<Utc>,
    pub to: DateTime<Utc>,
}

#[derive(Debug, Default, Clone)]
pub struct InMemoryEnvelopeStore {
    envelopes: Arc<RwLock<HashMap<String, Envelope>>>,
    uri_index: Arc<RwLock<HashMap<String, String>>>, // uri_hash -> envelope_id
    version_index: Arc<RwLock<HashMap<String, HashMap<String, Envelope>>>>, // envelope_id -> version -> envelope
    history_index: Arc<RwLock<HashMap<String, EnvelopeHistory>>>, // doc_id -> history
    deleted: Arc<RwLock<HashMap<String, DateTime<Utc>>>>, // envelope_id -> deletion_time
}

impl InMemoryEnvelopeStore {
    pub fn new() -> Self {
        InMemoryEnvelopeStore::default()
    }
    
    pub fn get_stats(&self) -> StoreStats {
        let envelopes = self.envelopes.read().unwrap();
        let deleted = self.deleted.read().unwrap();
        let version_index = self.version_index.read().unwrap();
        
        let total_versions: usize = version_index.values()
            .map(|versions| versions.len())
            .sum();
            
        StoreStats {
            total_envelopes: envelopes.len(),
            total_versions,
            deleted_envelopes: deleted.len(),
            memory_usage_estimate: self.estimate_memory_usage(),
        }
    }
    
    fn estimate_memory_usage(&self) -> usize {
        // Rough estimate of memory usage
        let envelopes = self.envelopes.read().unwrap();
        envelopes.values()
            .map(|env| env.content.raw.len() + 1024) // content + overhead estimate
            .sum()
    }
    
    fn matches_search_query(&self, envelope: &Envelope, query: &SearchQuery) -> bool {
        // Text search in content
        if let Some(ref text) = query.text {
            let content_lower = envelope.content.raw.to_lowercase();
            let title_lower = envelope.metadata.title
                .as_ref()
                .map(|t| t.to_lowercase())
                .unwrap_or_default();
            
            let text_lower = text.to_lowercase();
            if !content_lower.contains(&text_lower) && !title_lower.contains(&text_lower) {
                return false;
            }
        }
        
        // Tag matching
        if !query.tags.is_empty() {
            let envelope_tags: std::collections::HashSet<_> = envelope.metadata.tags.iter().collect();
            let query_tags: std::collections::HashSet<_> = query.tags.iter().collect();
            if envelope_tags.intersection(&query_tags).count() == 0 {
                return false;
            }
        }
        
        // Author matching
        if let Some(ref author) = query.author {
            if envelope.metadata.author.as_ref() != Some(author) {
                return false;
            }
        }
        
        // Content type matching
        if let Some(ref content_type) = query.content_type {
            if &envelope.content.content_type != content_type {
                return false;
            }
        }
        
        // Date range matching
        if let Some(ref date_range) = query.date_range {
            if envelope.timestamp < date_range.from || envelope.timestamp > date_range.to {
                return false;
            }
        }
        
        true
    }
}

#[async_trait::async_trait]
impl EnvelopeStore for InMemoryEnvelopeStore {
    async fn store(&self, envelope: &Envelope) -> Result<()> {
        let mut envelopes = self.envelopes.write().unwrap();
        let mut uri_index = self.uri_index.write().unwrap();
        let mut version_index = self.version_index.write().unwrap();
        let mut history_index = self.history_index.write().unwrap();
        
        // Store the envelope
        envelopes.insert(envelope.id.clone(), envelope.clone());
        
        // Update URI index
        let uri_hash = envelope.ucxl_uri.to_content_hash();
        uri_index.insert(uri_hash, envelope.id.clone());
        
        // Update version index
        version_index.entry(envelope.id.clone())
            .or_insert_with(HashMap::new)
            .insert(envelope.version.clone(), envelope.clone());
        
        // Update history
        let doc_id = envelope.ucxl_uri.path.clone(); // Use path as doc_id for now
        let version_info = VersionInfo {
            version_id: envelope.version.clone(),
            timestamp: envelope.timestamp,
            parent_versions: envelope.parent_version.iter().cloned().collect(),
            author: envelope.metadata.author.clone(),
            commit_message: None,
            content_hash: envelope.content_hash.clone(),
        };
        
        history_index.entry(doc_id.clone())
            .or_insert_with(|| EnvelopeHistory {
                doc_id,
                versions: Vec::new(),
                branches: HashMap::new(),
            })
            .versions.push(version_info);
        
        Ok(())
    }
    
    async fn retrieve(&self, envelope_id: &str) -> Result<Option<Envelope>> {
        let envelopes = self.envelopes.read().unwrap();
        let deleted = self.deleted.read().unwrap();
        
        if deleted.contains_key(envelope_id) {
            return Ok(None);
        }
        
        Ok(envelopes.get(envelope_id).cloned())
    }
    
    async fn retrieve_by_uri(&self, uri: &UCXLUri) -> Result<Option<Envelope>> {
        let uri_hash = uri.to_content_hash();
        
        let envelope_id = {
            let uri_index = self.uri_index.read().unwrap();
            uri_index.get(&uri_hash).cloned()
        };
        
        if let Some(envelope_id) = envelope_id {
            self.retrieve(&envelope_id).await
        } else {
            Ok(None)
        }
    }
    
    async fn retrieve_version(&self, envelope_id: &str, version: &str) -> Result<Option<Envelope>> {
        let version_index = self.version_index.read().unwrap();
        let deleted = self.deleted.read().unwrap();
        
        if deleted.contains_key(envelope_id) {
            return Ok(None);
        }
        
        if let Some(versions) = version_index.get(envelope_id) {
            Ok(versions.get(version).cloned())
        } else {
            Ok(None)
        }
    }
    
    async fn list_versions(&self, envelope_id: &str) -> Result<Vec<VersionInfo>> {
        let history_index = self.history_index.read().unwrap();
        
        // Find doc_id for this envelope_id
        for history in history_index.values() {
            if history.versions.iter().any(|v| v.version_id == envelope_id) {
                return Ok(history.versions.clone());
            }
        }
        
        Ok(Vec::new())
    }
    
    async fn get_history(&self, doc_id: &str) -> Result<EnvelopeHistory> {
        let history_index = self.history_index.read().unwrap();
        
        history_index.get(doc_id)
            .cloned()
            .ok_or_else(|| UCXLError::EnvelopeNotFound(doc_id.to_string()))
    }
    
    async fn delete(&self, envelope_id: &str, soft_delete: bool) -> Result<()> {
        if soft_delete {
            let mut deleted = self.deleted.write().unwrap();
            deleted.insert(envelope_id.to_string(), Utc::now());
        } else {
            let mut envelopes = self.envelopes.write().unwrap();
            let mut version_index = self.version_index.write().unwrap();
            
            envelopes.remove(envelope_id);
            version_index.remove(envelope_id);
        }
        
        Ok(())
    }
    
    async fn search(&self, query: &SearchQuery) -> Result<Vec<Envelope>> {
        let envelopes = self.envelopes.read().unwrap();
        let deleted = self.deleted.read().unwrap();
        
        let mut results: Vec<_> = envelopes.values()
            .filter(|env| !deleted.contains_key(&env.id))
            .filter(|env| self.matches_search_query(env, query))
            .cloned()
            .collect();
        
        // Sort by timestamp (newest first)
        results.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
        
        // Apply pagination
        if let Some(offset) = query.offset {
            if offset < results.len() {
                results = results[offset..].to_vec();
            } else {
                results.clear();
            }
        }
        
        if let Some(limit) = query.limit {
            results.truncate(limit);
        }
        
        Ok(results)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoreStats {
    pub total_envelopes: usize,
    pub total_versions: usize,
    pub deleted_envelopes: usize,
    pub memory_usage_estimate: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{UCXLUri, envelope::{EnvelopeMetadata}};
    
    #[tokio::test]
    async fn test_in_memory_store_basic_operations() {
        let store = InMemoryEnvelopeStore::new();
        let uri = UCXLUri::new("ucxl://example.com/test").unwrap();
        
        let metadata = EnvelopeMetadata {
            author: Some("test_author".to_string()),
            title: Some("Test Document".to_string()),
            tags: vec!["test".to_string()],
            source: None,
            context_data: HashMap::new(),
        };
        
        let envelope = crate::Envelope::new(
            uri.clone(),
            "# Test Content".to_string(),
            "text/markdown".to_string(),
            metadata,
        ).unwrap();
        
        // Store envelope
        store.store(&envelope).await.unwrap();
        
        // Retrieve by ID
        let retrieved = store.retrieve(&envelope.id).await.unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().id, envelope.id);
        
        // Retrieve by URI
        let retrieved_by_uri = store.retrieve_by_uri(&uri).await.unwrap();
        assert!(retrieved_by_uri.is_some());
        assert_eq!(retrieved_by_uri.unwrap().id, envelope.id);
    }
    
    #[tokio::test]
    async fn test_search_functionality() {
        let store = InMemoryEnvelopeStore::new();
        
        // Store multiple envelopes
        for i in 0..5 {
            let uri = UCXLUri::new(&format!("ucxl://example.com/doc{}", i)).unwrap();
            let metadata = EnvelopeMetadata {
                author: Some(format!("author_{}", i)),
                title: Some(format!("Document {}", i)),
                tags: vec![format!("tag_{}", i % 2)],
                source: None,
                context_data: HashMap::new(),
            };
            
            let envelope = crate::Envelope::new(
                uri,
                format!("Content for document {}", i),
                "text/markdown".to_string(),
                metadata,
            ).unwrap();
            
            store.store(&envelope).await.unwrap();
        }
        
        // Search by tag
        let query = SearchQuery {
            text: None,
            tags: vec!["tag_0".to_string()],
            author: None,
            content_type: None,
            date_range: None,
            limit: None,
            offset: None,
        };
        
        let results = store.search(&query).await.unwrap();
        assert_eq!(results.len(), 3); // Documents 0, 2, 4
        
        // Search by author
        let query = SearchQuery {
            text: None,
            tags: vec![],
            author: Some("author_1".to_string()),
            content_type: None,
            date_range: None,
            limit: None,
            offset: None,
        };
        
        let results = store.search(&query).await.unwrap();
        assert_eq!(results.len(), 1);
    }
}