use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use serde::{Serialize, Deserialize};
use sha2::{Digest, Sha256};
use chrono::{DateTime, Utc};
// use futures::stream::StreamExt;
use libp2p::{
    identity, Multiaddr, PeerId,
};
use crate::{Envelope, Result, UCXLError, EnvelopeStore, SearchQuery, VersionInfo, EnvelopeHistory};

// Serde helpers for PeerId
mod peer_id_serde {
    use super::*;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    
    pub fn serialize<S>(peer_id: &PeerId, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        peer_id.to_string().serialize(serializer)
    }
    
    pub fn deserialize<'de, D>(deserializer: D) -> std::result::Result<PeerId, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        s.parse().map_err(serde::de::Error::custom)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DHTRecord {
    pub key: String,
    pub value: Vec<u8>,
    pub timestamp: DateTime<Utc>,
    pub content_hash: String,
    pub metadata: DHTMetadata,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DHTMetadata {
    pub content_type: String,
    pub size: usize,
    pub ttl: Option<u64>, // TTL in seconds
    pub replicas: u8,     // Number of replicas to maintain
    pub publisher: String, // Node ID that published this record
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DHTNode {
    #[serde(with = "peer_id_serde")]
    pub peer_id: PeerId,
    pub addresses: Vec<Multiaddr>,
    pub last_seen: DateTime<Utc>,
    pub capabilities: NodeCapabilities,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeCapabilities {
    pub storage_capacity: u64,
    pub bandwidth_limit: u64,
    pub supports_versioning: bool,
    pub supports_search: bool,
}

#[derive(Debug)]
pub struct DHTEnvelopeStore {
    // Local cache for fast access
    local_cache: Arc<RwLock<HashMap<String, Envelope>>>,
    dht_records: Arc<RwLock<HashMap<String, DHTRecord>>>,
    
    // DHT network state
    local_peer_id: PeerId,
    known_peers: Arc<RwLock<HashMap<PeerId, DHTNode>>>,
    
    // Configuration
    config: DHTConfig,
    
    // Statistics
    stats: Arc<RwLock<DHTStats>>,
}

#[derive(Debug, Clone)]
pub struct DHTConfig {
    pub replication_factor: u8,
    pub cache_size_limit: usize,
    pub ttl_seconds: u64,
    pub bootstrap_nodes: Vec<Multiaddr>,
    pub enable_content_routing: bool,
    pub enable_peer_discovery: bool,
}

#[derive(Debug, Default)]
pub struct DHTStats {
    pub total_records: usize,
    pub cached_records: usize,
    pub network_puts: u64,
    pub network_gets: u64,
    pub cache_hits: u64,
    pub cache_misses: u64,
    pub replication_operations: u64,
    pub connected_peers: usize,
    pub storage_used: u64,
}

impl Default for DHTConfig {
    fn default() -> Self {
        DHTConfig {
            replication_factor: 3,
            cache_size_limit: 10000,
            ttl_seconds: 86400, // 24 hours
            bootstrap_nodes: Vec::new(),
            enable_content_routing: true,
            enable_peer_discovery: true,
        }
    }
}

impl DHTEnvelopeStore {
    pub fn new(config: DHTConfig) -> Self {
        let local_key = identity::Keypair::generate_ed25519();
        let local_peer_id = PeerId::from(local_key.public());
        
        DHTEnvelopeStore {
            local_cache: Arc::new(RwLock::new(HashMap::new())),
            dht_records: Arc::new(RwLock::new(HashMap::new())),
            local_peer_id,
            known_peers: Arc::new(RwLock::new(HashMap::new())),
            config,
            stats: Arc::new(RwLock::new(DHTStats::default())),
        }
    }
    
    pub fn with_bootstrap_nodes(mut self, nodes: Vec<Multiaddr>) -> Self {
        self.config.bootstrap_nodes = nodes;
        self
    }
    
    pub async fn start_network(&self) -> Result<()> {
        // Initialize the DHT network
        tracing::info!("Starting DHT network for peer: {}", self.local_peer_id);
        
        // Bootstrap from known nodes
        for addr in &self.config.bootstrap_nodes {
            self.connect_to_peer(addr.clone()).await?;
        }
        
        Ok(())
    }
    
    async fn connect_to_peer(&self, addr: Multiaddr) -> Result<()> {
        tracing::debug!("Attempting to connect to peer: {}", addr);
        // TODO: Implement actual peer connection logic
        Ok(())
    }
    
    fn compute_dht_key(&self, envelope_id: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(format!("envelope:{}", envelope_id));
        hex::encode(hasher.finalize())
    }
    
    fn compute_uri_key(&self, uri_hash: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(format!("uri:{}", uri_hash));
        hex::encode(hasher.finalize())
    }
    
    async fn store_in_dht(&self, key: &str, envelope: &Envelope) -> Result<()> {
        let serialized = bincode::serialize(envelope)
            .map_err(|e| UCXLError::SerializationError(e.to_string()))?;
        
        let dht_record = DHTRecord {
            key: key.to_string(),
            value: serialized,
            timestamp: Utc::now(),
            content_hash: envelope.content_hash.clone(),
            metadata: DHTMetadata {
                content_type: envelope.content.content_type.clone(),
                size: envelope.content.raw.len(),
                ttl: Some(self.config.ttl_seconds),
                replicas: self.config.replication_factor,
                publisher: self.local_peer_id.to_string(),
            },
        };
        
        // Store in local DHT records
        {
            let mut records = self.dht_records.write().unwrap();
            records.insert(key.to_string(), dht_record.clone());
        }
        
        // Update statistics
        {
            let mut stats = self.stats.write().unwrap();
            stats.network_puts += 1;
            stats.total_records += 1;
            stats.storage_used += envelope.content.raw.len() as u64;
        }
        
        // TODO: Implement actual DHT put operation using libp2p-kad
        tracing::debug!("Stored envelope {} in DHT with key {}", envelope.id, key);
        
        Ok(())
    }
    
    async fn retrieve_from_dht(&self, key: &str) -> Result<Option<Envelope>> {
        // First check local cache
        {
            let cache = self.local_cache.read().unwrap();
            if let Some(envelope) = cache.get(key) {
                let mut stats = self.stats.write().unwrap();
                stats.cache_hits += 1;
                return Ok(Some(envelope.clone()));
            }
        }
        
        // Check DHT records
        let dht_record = {
            let records = self.dht_records.read().unwrap();
            records.get(key).cloned()
        };
        
        if let Some(record) = dht_record {
            let envelope: Envelope = bincode::deserialize(&record.value)
                .map_err(|e| UCXLError::SerializationError(e.to_string()))?;
            
            // Cache the result
            {
                let mut cache = self.local_cache.write().unwrap();
                if cache.len() >= self.config.cache_size_limit {
                    // Simple LRU: remove first item
                    if let Some(first_key) = cache.keys().next().cloned() {
                        cache.remove(&first_key);
                    }
                }
                cache.insert(key.to_string(), envelope.clone());
            }
            
            // Update statistics
            {
                let mut stats = self.stats.write().unwrap();
                stats.network_gets += 1;
                stats.cached_records += 1;
            }
            
            Ok(Some(envelope))
        } else {
            let mut stats = self.stats.write().unwrap();
            stats.cache_misses += 1;
            Ok(None)
        }
    }
    
    pub fn get_dht_stats(&self) -> DHTStats {
        let stats = self.stats.read().unwrap();
        DHTStats {
            total_records: stats.total_records,
            cached_records: self.local_cache.read().unwrap().len(),
            network_puts: stats.network_puts,
            network_gets: stats.network_gets,
            cache_hits: stats.cache_hits,
            cache_misses: stats.cache_misses,
            replication_operations: stats.replication_operations,
            connected_peers: self.known_peers.read().unwrap().len(),
            storage_used: stats.storage_used,
        }
    }
    
    pub async fn replicate_content(&self, key: &str, target_replicas: u8) -> Result<u8> {
        // TODO: Implement content replication across multiple peers
        tracing::debug!("Replicating content for key: {} to {} peers", key, target_replicas);
        
        let mut stats = self.stats.write().unwrap();
        stats.replication_operations += 1;
        
        // For now, assume successful replication
        Ok(target_replicas)
    }
    
    pub async fn discover_peers(&self) -> Result<Vec<DHTNode>> {
        // TODO: Implement peer discovery using libp2p-kad
        let peers = self.known_peers.read().unwrap();
        Ok(peers.values().cloned().collect())
    }
    
    async fn cleanup_expired_records(&self) -> Result<usize> {
        let now = Utc::now();
        let expired_count;
        
        {
            let mut records = self.dht_records.write().unwrap();
            let initial_count = records.len();
            
            records.retain(|_key, record| {
                if let Some(ttl) = record.metadata.ttl {
                    let age = now.timestamp() as u64 - record.timestamp.timestamp() as u64;
                    age < ttl
                } else {
                    true // Keep records without TTL
                }
            });
            
            expired_count = initial_count - records.len();
        }
        
        if expired_count > 0 {
            let mut stats = self.stats.write().unwrap();
            stats.total_records = stats.total_records.saturating_sub(expired_count);
            tracing::debug!("Cleaned up {} expired DHT records", expired_count);
        }
        
        Ok(expired_count)
    }
}

#[async_trait::async_trait]
impl EnvelopeStore for DHTEnvelopeStore {
    async fn store(&self, envelope: &Envelope) -> Result<()> {
        // Store in DHT with envelope ID as key
        let dht_key = self.compute_dht_key(&envelope.id);
        self.store_in_dht(&dht_key, envelope).await?;
        
        // Also store URI mapping
        let uri_key = self.compute_uri_key(&envelope.ucxl_uri.to_content_hash());
        self.store_in_dht(&uri_key, envelope).await?;
        
        // Cache locally for fast access
        {
            let mut cache = self.local_cache.write().unwrap();
            cache.insert(envelope.id.clone(), envelope.clone());
        }
        
        tracing::debug!("Stored envelope {} in DHT and local cache", envelope.id);
        Ok(())
    }
    
    async fn retrieve(&self, envelope_id: &str) -> Result<Option<Envelope>> {
        let dht_key = self.compute_dht_key(envelope_id);
        self.retrieve_from_dht(&dht_key).await
    }
    
    async fn retrieve_by_uri(&self, uri: &crate::UCXLUri) -> Result<Option<Envelope>> {
        let uri_hash = uri.to_content_hash();
        let uri_key = self.compute_uri_key(&uri_hash);
        self.retrieve_from_dht(&uri_key).await
    }
    
    async fn retrieve_version(&self, envelope_id: &str, version: &str) -> Result<Option<Envelope>> {
        // For now, use the same retrieval mechanism
        // TODO: Implement version-specific retrieval
        let version_key = format!("{}:{}", envelope_id, version);
        let dht_key = self.compute_dht_key(&version_key);
        self.retrieve_from_dht(&dht_key).await
    }
    
    async fn list_versions(&self, _envelope_id: &str) -> Result<Vec<VersionInfo>> {
        // TODO: Implement version listing from DHT
        // For now, return empty list
        Ok(Vec::new())
    }
    
    async fn get_history(&self, doc_id: &str) -> Result<EnvelopeHistory> {
        // TODO: Implement history retrieval from DHT
        // For now, return empty history
        Ok(EnvelopeHistory {
            doc_id: doc_id.to_string(),
            versions: Vec::new(),
            branches: HashMap::new(),
        })
    }
    
    async fn delete(&self, envelope_id: &str, soft_delete: bool) -> Result<()> {
        let dht_key = self.compute_dht_key(envelope_id);
        
        if soft_delete {
            // Mark as deleted in DHT metadata
            // TODO: Implement soft delete marking
            tracing::debug!("Soft deleting envelope {} from DHT", envelope_id);
        } else {
            // Remove from DHT and cache
            {
                let mut records = self.dht_records.write().unwrap();
                records.remove(&dht_key);
            }
            {
                let mut cache = self.local_cache.write().unwrap();
                cache.remove(envelope_id);
            }
            
            let mut stats = self.stats.write().unwrap();
            stats.total_records = stats.total_records.saturating_sub(1);
            stats.cached_records = stats.cached_records.saturating_sub(1);
            
            tracing::debug!("Hard deleted envelope {} from DHT and cache", envelope_id);
        }
        
        Ok(())
    }
    
    async fn search(&self, query: &SearchQuery) -> Result<Vec<Envelope>> {
        // TODO: Implement distributed search across DHT
        // For now, search local cache and DHT records
        let mut results = Vec::new();
        
        // Search cached envelopes
        {
            let cache = self.local_cache.read().unwrap();
            for envelope in cache.values() {
                if self.matches_query(envelope, query) {
                    results.push(envelope.clone());
                }
            }
        }
        
        // Search DHT records
        {
            let records = self.dht_records.read().unwrap();
            for record in records.values() {
                if let Ok(envelope) = bincode::deserialize::<Envelope>(&record.value) {
                    if self.matches_query(&envelope, query) && !results.iter().any(|e| e.id == envelope.id) {
                        results.push(envelope);
                    }
                }
            }
        }
        
        // 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)
    }
}

impl DHTEnvelopeStore {
    fn matches_query(&self, envelope: &Envelope, query: &SearchQuery) -> bool {
        // Text search
        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
    }
}

// Background task for DHT maintenance
pub struct DHTMaintenance {
    store: Arc<DHTEnvelopeStore>,
    cleanup_interval: Duration,
    replication_interval: Duration,
}

impl DHTMaintenance {
    pub fn new(store: Arc<DHTEnvelopeStore>) -> Self {
        DHTMaintenance {
            store,
            cleanup_interval: Duration::from_secs(300), // 5 minutes
            replication_interval: Duration::from_secs(600), // 10 minutes
        }
    }
    
    pub async fn start_maintenance_loop(&self) {
        let mut cleanup_timer = tokio::time::interval(self.cleanup_interval);
        let mut replication_timer = tokio::time::interval(self.replication_interval);
        
        loop {
            tokio::select! {
                _ = cleanup_timer.tick() => {
                    if let Err(e) = self.store.cleanup_expired_records().await {
                        tracing::warn!("DHT cleanup failed: {}", e);
                    }
                }
                _ = replication_timer.tick() => {
                    // TODO: Check replication health and repair if needed
                    tracing::debug!("DHT replication check completed");
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{UCXLUri, envelope::{EnvelopeMetadata}};
    use std::collections::HashMap;
    
    #[tokio::test]
    async fn test_dht_store_basic_operations() {
        let config = DHTConfig::default();
        let store = DHTEnvelopeStore::new(config);
        
        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,
            "# 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);
        
        // Check DHT stats
        let stats = store.get_dht_stats();
        assert_eq!(stats.total_records, 2); // One for envelope ID, one for URI mapping
        assert!(stats.storage_used > 0);
    }
    
    #[tokio::test]
    async fn test_dht_search() {
        let config = DHTConfig::default();
        let store = DHTEnvelopeStore::new(config);
        
        // Store test envelopes
        for i in 0..3 {
            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(), 2); // Documents 0 and 2
    }
}