//! Caching identity resolver implementation.
//!
//! Provides a three-layer caching strategy for DID document resolution:
//! 1. In-memory LRU cache (fastest)
//! 2. Database storage (persistent)  
//! 3. Base resolver (network fetch)

use anyhow::Result;
use async_trait::async_trait;
use atproto_identity::{model::Document, resolve::IdentityResolver, traits::DidDocumentStorage};
use chrono::{Duration, Utc};

// TODO: Use a different library because lru uses unsafe.
use lru::LruCache;

use std::num::NonZeroUsize;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::warn;

/// Configuration for the caching identity resolver
#[derive(Clone, Debug)]
pub struct CacheConfig {
    /// Maximum number of entries in the in-memory cache
    pub memory_cache_size: usize,
    /// TTL for in-memory cache entries (in seconds)
    pub memory_ttl_seconds: i64,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            memory_cache_size: 1000,
            memory_ttl_seconds: 300,
        }
    }
}

/// Entry in the memory cache with timestamp
#[derive(Clone)]
struct CachedDocument {
    document: Document,
    cached_at: chrono::DateTime<Utc>,
}

impl CachedDocument {
    fn is_expired(&self, ttl_seconds: i64) -> bool {
        let age = Utc::now() - self.cached_at;
        age > Duration::seconds(ttl_seconds)
    }
}

/// A caching identity resolver that uses multiple layers of caching.
///
/// Resolution order:
/// 1. Check in-memory LRU cache
/// 2. Check database storage
/// 3. Resolve using base resolver
/// 4. Update both caches with the result
pub struct CachingIdentityResolver<R, S>
where
    R: IdentityResolver + 'static,
    S: DidDocumentStorage + 'static,
{
    /// The base resolver for actual DID resolution
    base_resolver: Arc<R>,
    /// Database storage for persistent caching
    storage: Arc<S>,
    /// In-memory LRU cache
    memory_cache: Arc<RwLock<LruCache<String, CachedDocument>>>,

    /// Cache configuration
    config: CacheConfig,
}

impl<R, S> CachingIdentityResolver<R, S>
where
    R: IdentityResolver + 'static,
    S: DidDocumentStorage + 'static,
{
    /// Creates a new caching identity resolver with default configuration
    pub fn new(base_resolver: Arc<R>, storage: Arc<S>) -> Self {
        Self::with_config(base_resolver, storage, CacheConfig::default())
    }

    /// Creates a new caching identity resolver with custom configuration
    pub fn with_config(base_resolver: Arc<R>, storage: Arc<S>, config: CacheConfig) -> Self {
        let cache_size = NonZeroUsize::new(config.memory_cache_size.max(1))
            .expect("Cache size must be at least 1");

        Self {
            base_resolver,
            storage,
            memory_cache: Arc::new(RwLock::new(LruCache::new(cache_size))),
            config,
        }
    }

    /// Normalize the subject to a consistent format for caching
    fn normalize_subject(subject: &str) -> String {
        // Convert handles to lowercase, keep DIDs as-is
        if subject.starts_with("did:") {
            subject.to_string()
        } else {
            subject.to_lowercase()
        }
    }

    /// Try to get a document from the in-memory cache
    async fn get_from_memory(&self, subject: &str) -> Option<Document> {
        let normalized = Self::normalize_subject(subject);
        let mut cache = self.memory_cache.write().await;

        if let Some(cached) = cache.get(&normalized) {
            if !cached.is_expired(self.config.memory_ttl_seconds) {
                return Some(cached.document.clone());
            } else {
                // Remove expired entry
                cache.pop(&normalized);
            }
        }

        None
    }

    /// Try to get a document from database storage
    async fn get_from_storage(&self, subject: &str) -> Option<Document> {
        // First resolve the subject to a DID if it's a handle
        // This is tricky because we need the DID to query storage
        // For now, we'll only check storage if the subject is already a DID
        if !subject.starts_with("did:") {
            return None;
        }

        match self.storage.as_ref().get_document_by_did(subject).await {
            Ok(Some(document)) => {
                // Check if the database entry is still fresh enough
                // Note: The storage implementation doesn't provide timestamps,
                // so we'll trust it for now. In a real implementation, you'd
                // want to add timestamp tracking to the storage layer.
                Some(document)
            }
            Ok(None) => None,
            Err(_) => {
                warn!("Failed to query database cache");
                None
            }
        }
    }

    /// Store a document in both memory and database caches
    async fn store_in_caches(&self, subject: &str, document: Document) {
        let normalized = Self::normalize_subject(subject);

        // Store in memory cache
        {
            let mut cache = self.memory_cache.write().await;
            cache.put(
                normalized.clone(),
                CachedDocument {
                    document: document.clone(),
                    cached_at: Utc::now(),
                },
            );
        }

        // Store in database
        if let Err(e) = self.storage.as_ref().store_document(document.clone()).await {
            warn!("Failed to store document in database cache: {}", e);
        }

        // Also store by handle if the subject was a handle
        if !subject.starts_with("did:") {
            // Store the handle -> document mapping in memory
            let mut cache = self.memory_cache.write().await;
            cache.put(
                Self::normalize_subject(subject),
                CachedDocument {
                    document: document.clone(),
                    cached_at: Utc::now(),
                },
            );
        }
    }

    /// Refresh a stale entry in the background
    #[allow(dead_code)]
    async fn background_refresh(&self, subject: String) {
        let resolver = self.base_resolver.clone();
        let storage = self.storage.clone();
        let memory_cache = self.memory_cache.clone();

        tokio::spawn(async move {
            match resolver.resolve(&subject).await {
                Ok(document) => {
                    // Update memory cache
                    let normalized = Self::normalize_subject(&subject);
                    let mut cache = memory_cache.write().await;
                    cache.put(
                        normalized,
                        CachedDocument {
                            document: document.clone(),
                            cached_at: Utc::now(),
                        },
                    );

                    // Update database
                    if let Err(e) = storage.as_ref().store_document(document).await {
                        warn!("Failed to update database during background refresh: {}", e);
                    }
                }
                Err(_) => {
                    warn!("Background refresh failed");
                }
            }
        });
    }

    /// Clear all caches
    pub async fn clear_caches(&self) -> Result<()> {
        // Clear memory cache
        let mut cache = self.memory_cache.write().await;
        cache.clear();

        // Note: Database clearing would need to be implemented in the storage layer
        // For now, we just clear the memory cache

        Ok(())
    }

    /// Get cache statistics
    pub async fn cache_stats(&self) -> CacheStats {
        let cache = self.memory_cache.read().await;
        CacheStats {
            memory_entries: cache.len(),
            memory_capacity: cache.cap().get(),
        }
    }
}

#[async_trait]
impl<R, S> IdentityResolver for CachingIdentityResolver<R, S>
where
    R: IdentityResolver + 'static,
    S: DidDocumentStorage + 'static,
{
    async fn resolve(&self, subject: &str) -> Result<Document> {
        // 1. Check memory cache
        if let Some(document) = self.get_from_memory(subject).await {
            return Ok(document);
        }

        // 2. Check database storage
        if let Some(document) = self.get_from_storage(subject).await {
            // Store in memory cache for faster future access
            self.store_in_caches(subject, document.clone()).await;
            return Ok(document);
        }

        // 3. Resolve using base resolver
        let document = self.base_resolver.resolve(subject).await?;

        // 4. Store in both caches
        self.store_in_caches(subject, document.clone()).await;

        Ok(document)
    }
}

/// Statistics about the cache
#[derive(Debug, Clone)]
pub struct CacheStats {
    pub memory_entries: usize,
    pub memory_capacity: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    /// Mock storage for testing
    struct MockStorage {
        documents: Arc<RwLock<HashMap<String, Document>>>,
    }

    impl MockStorage {
        fn new() -> Self {
            Self {
                documents: Arc::new(RwLock::new(HashMap::new())),
            }
        }
    }

    #[async_trait]
    impl DidDocumentStorage for MockStorage {
        async fn get_document_by_did(&self, did: &str) -> Result<Option<Document>> {
            let docs = self.documents.read().await;
            Ok(docs.get(did).cloned())
        }

        async fn store_document(&self, document: Document) -> Result<()> {
            let mut docs = self.documents.write().await;
            docs.insert(document.id.clone(), document);
            Ok(())
        }

        async fn delete_document_by_did(&self, did: &str) -> Result<()> {
            let mut docs = self.documents.write().await;
            docs.remove(did);
            Ok(())
        }
    }

    /// Mock resolver for testing
    struct MockResolver {
        documents: Arc<RwLock<HashMap<String, Document>>>,
        call_count: Arc<RwLock<usize>>,
    }

    impl MockResolver {
        fn new() -> Self {
            Self {
                documents: Arc::new(RwLock::new(HashMap::new())),
                call_count: Arc::new(RwLock::new(0)),
            }
        }

        async fn add_document(&self, did: String, document: Document) {
            let mut docs = self.documents.write().await;
            docs.insert(did, document);
        }

        async fn get_call_count(&self) -> usize {
            *self.call_count.read().await
        }
    }

    #[async_trait]
    impl IdentityResolver for MockResolver {
        async fn resolve(&self, subject: &str) -> Result<Document> {
            let mut count = self.call_count.write().await;
            *count += 1;

            let docs = self.documents.read().await;
            docs.get(subject)
                .cloned()
                .ok_or_else(|| anyhow::anyhow!("Document not found"))
        }
    }

    #[tokio::test]
    async fn test_caching_resolver_memory_cache() {
        let base_resolver = Arc::new(MockResolver::new());
        let storage = Arc::new(MockStorage::new());

        let test_did = "did:plc:test123";
        let test_doc = Document {
            id: test_did.to_string(),
            also_known_as: vec!["test.bsky.social".to_string()],
            verification_method: vec![],
            service: vec![],
            context: vec![],
            extra: Default::default(),
        };

        base_resolver
            .add_document(test_did.to_string(), test_doc.clone())
            .await;

        let config = CacheConfig {
            memory_cache_size: 10,
            memory_ttl_seconds: 60,
        };

        let caching_resolver =
            CachingIdentityResolver::with_config(base_resolver.clone(), storage, config);

        // First resolve should hit the base resolver
        let result1 = caching_resolver.resolve(test_did).await.unwrap();
        assert_eq!(result1.id, test_did.to_string());
        assert_eq!(base_resolver.get_call_count().await, 1);

        // Second resolve should hit the memory cache
        let result2 = caching_resolver.resolve(test_did).await.unwrap();
        assert_eq!(result2.id, test_did.to_string());
        assert_eq!(base_resolver.get_call_count().await, 1); // No additional call

        // Verify cache stats
        let stats = caching_resolver.cache_stats().await;
        assert_eq!(stats.memory_entries, 1);
        assert_eq!(stats.memory_capacity, 10);
    }

    #[tokio::test]
    async fn test_handle_normalization() {
        let base_resolver = Arc::new(MockResolver::new());
        let storage = Arc::new(MockStorage::new());

        let test_handle = "Test.BSKY.Social";
        let normalized_handle = "test.bsky.social";
        let test_doc = Document {
            id: "did:plc:test123".to_string(),
            also_known_as: vec![normalized_handle.to_string()],
            verification_method: vec![],
            service: vec![],
            context: vec![],
            extra: Default::default(),
        };

        base_resolver
            .add_document(test_handle.to_string(), test_doc.clone())
            .await;
        base_resolver
            .add_document(normalized_handle.to_string(), test_doc.clone())
            .await;

        let caching_resolver = CachingIdentityResolver::new(base_resolver.clone(), storage);

        // Resolve with uppercase handle
        let result1 = caching_resolver.resolve(test_handle).await.unwrap();
        assert_eq!(base_resolver.get_call_count().await, 1);

        // Resolve with lowercase handle - should hit cache
        let result2 = caching_resolver.resolve(normalized_handle).await.unwrap();
        assert_eq!(base_resolver.get_call_count().await, 1); // No additional call

        assert_eq!(result1.id, result2.id);
    }
}