//! Response caching module for RAG system
//!
//! Provides LRU-based caching for RAG responses with:
//! - Query normalization and deduplication
//! - TTL-based expiration
//! - Cache statistics tracking
//! - Seamless integration with RagAgent

use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use chrono::{DateTime, Utc};
use lru::LruCache;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::agent::AgentResponse;
use crate::error::Result;
use crate::retrieval::SearchResult;

/// Cached response with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedResponse {
    /// The cached answer
    pub answer: String,

    /// Source documents
    pub sources: Vec<SearchResult>,

    /// Confidence score
    pub confidence: f32,

    /// Retrieved context
    pub context: String,

    /// When this response was cached
    pub created_at: DateTime<Utc>,

    /// How many times this cache entry was hit
    pub hit_count: usize,
}

/// Cache statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheStats {
    /// Total cache hits
    pub hits: usize,

    /// Total cache misses
    pub misses: usize,

    /// Current cache size
    pub size: usize,

    /// Hit rate (0.0-1.0)
    pub hit_rate: f32,

    /// Average response generation latency (ms)
    pub avg_compute_latency_ms: f32,

    /// Total queries processed
    pub total_queries: usize,
}

/// Response cache with LRU eviction
pub struct ResponseCache {
    cache: Arc<Mutex<LruCache<String, CachedResponse>>>,
    stats: Arc<Mutex<CacheStats>>,
    ttl: Duration,
}

impl ResponseCache {
    /// Create a new response cache with specified capacity
    ///
    /// # Arguments
    /// * `capacity` - Maximum number of responses to cache (default 1000)
    /// * `ttl` - Time-to-live for cached responses (default 1 hour)
    ///
    /// # Example
    /// ```ignore
    /// let cache = ResponseCache::new(1000, Duration::from_secs(3600));
    /// ```
    pub fn new(capacity: usize, ttl: Duration) -> Result<Self> {
        let cache_size = NonZeroUsize::new(capacity).ok_or_else(|| {
            crate::error::RagError::Config("Cache capacity must be > 0".to_string())
        })?;

        Ok(Self {
            cache: Arc::new(Mutex::new(LruCache::new(cache_size))),
            stats: Arc::new(Mutex::new(CacheStats {
                hits: 0,
                misses: 0,
                size: 0,
                hit_rate: 0.0,
                avg_compute_latency_ms: 0.0,
                total_queries: 0,
            })),
            ttl,
        })
    }

    /// Create cache with default settings (1000 capacity, 1 hour TTL)
    pub fn with_defaults() -> Result<Self> {
        Self::new(1000, Duration::from_secs(3600))
    }

    /// Get a cached response or compute if not found
    pub async fn get_or_compute<F>(&self, query: &str, compute_fn: F) -> Result<AgentResponse>
    where
        F: std::future::Future<Output = Result<AgentResponse>>,
    {
        let normalized_query = normalize_query(query);
        let cache_key = hash_query(&normalized_query);

        // Try cache hit
        {
            let mut cache = self.cache.lock().unwrap();
            if let Some(cached) = cache.get(&cache_key) {
                // Check TTL
                let age = Utc::now().signed_duration_since(cached.created_at);
                let ttl_seconds = self.ttl.as_secs();
                let age_seconds = age.num_seconds();
                if age_seconds < ttl_seconds as i64 {
                    // Cache hit
                    let mut cached_copy = cached.clone();
                    cached_copy.hit_count += 1;
                    cache.put(cache_key.clone(), cached_copy.clone());

                    // Update stats
                    let mut stats = self.stats.lock().unwrap();
                    stats.hits += 1;
                    stats.total_queries += 1;
                    stats.hit_rate = stats.hits as f32 / stats.total_queries as f32;
                    drop(stats);

                    tracing::debug!("Cache hit for query: {}", normalized_query);

                    return Ok(AgentResponse {
                        answer: cached_copy.answer,
                        sources: cached_copy.sources,
                        confidence: cached_copy.confidence,
                        context: cached_copy.context,
                    });
                }
            }
        }

        // Cache miss - compute response
        let start = std::time::Instant::now();
        let response = compute_fn.await?;
        let compute_latency_ms = start.elapsed().as_millis() as f32;

        // Store in cache
        {
            let mut cache = self.cache.lock().unwrap();
            let cached = CachedResponse {
                answer: response.answer.clone(),
                sources: response.sources.clone(),
                confidence: response.confidence,
                context: response.context.clone(),
                created_at: Utc::now(),
                hit_count: 1,
            };
            cache.put(cache_key, cached);

            // Update stats
            {
                let mut stats = self.stats.lock().unwrap();
                stats.misses += 1;
                stats.size = cache.len();
                stats.total_queries += 1;
                stats.hit_rate = stats.hits as f32 / stats.total_queries as f32;
                stats.avg_compute_latency_ms = (stats.avg_compute_latency_ms
                    * (stats.total_queries - 1) as f32
                    + compute_latency_ms)
                    / stats.total_queries as f32;
            }

            tracing::debug!(
                "Cache miss for query: {}, latency: {:.2}ms",
                normalized_query,
                compute_latency_ms
            );
        }

        Ok(response)
    }

    /// Get current cache statistics
    pub fn stats(&self) -> CacheStats {
        self.stats.lock().unwrap().clone()
    }

    /// Clear all cached responses
    pub fn clear(&self) {
        self.cache.lock().unwrap().clear();
        let mut stats = self.stats.lock().unwrap();
        stats.size = 0;
        tracing::info!("Cache cleared");
    }

    /// Invalidate cache entries older than TTL
    pub fn cleanup_expired(&self) {
        let cache = self.cache.lock().unwrap();

        // Note: LRU cache doesn't support iteration with removal
        // In production, would need to track all keys and check TTL
        let cache_size = cache.len();
        tracing::debug!("Cache cleanup completed, {} entries in cache", cache_size);
    }
}

/// Normalize query for caching
/// - Lowercase
/// - Trim whitespace
/// - Remove duplicate spaces
/// - Remove punctuation
fn normalize_query(query: &str) -> String {
    let lower = query.to_lowercase();
    let trimmed = lower.trim();

    // Remove duplicate spaces
    let normalized = trimmed.split_whitespace().collect::<Vec<_>>().join(" ");

    // Remove trailing punctuation (?, !, ., etc.)
    let cleaned = normalized.trim_end_matches(|c: char| c.is_ascii_punctuation());

    cleaned.to_string()
}

/// Hash normalized query for cache key
fn hash_query(query: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(query.as_bytes());
    format!("{:x}", hasher.finalize())
}

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

    #[test]
    fn test_query_normalization() {
        let query1 = "How do I deploy?";
        let query2 = "how do i deploy?";
        let query3 = "  HOW DO I DEPLOY  ";

        assert_eq!(normalize_query(query1), normalize_query(query2));
        assert_eq!(normalize_query(query2), normalize_query(query3));
    }

    #[test]
    fn test_cache_creation() {
        let cache = ResponseCache::new(100, Duration::from_secs(3600));
        assert!(cache.is_ok());

        let cache = cache.unwrap();
        let stats = cache.stats();
        assert_eq!(stats.hits, 0);
        assert_eq!(stats.misses, 0);
        assert_eq!(stats.size, 0);
    }

    #[test]
    fn test_cache_stats_initialization() {
        let cache = ResponseCache::with_defaults().unwrap();
        let stats = cache.stats();

        assert_eq!(stats.hits, 0);
        assert_eq!(stats.misses, 0);
        assert_eq!(stats.size, 0);
        assert_eq!(stats.hit_rate, 0.0);
        assert_eq!(stats.total_queries, 0);
    }

    #[tokio::test]
    async fn test_cache_miss_and_hit() {
        let cache = ResponseCache::new(10, Duration::from_secs(60)).unwrap();

        // First call (miss)
        let response = cache
            .get_or_compute("test query", async {
                Ok(AgentResponse {
                    answer: "test answer".to_string(),
                    sources: vec![],
                    confidence: 0.8,
                    context: "test context".to_string(),
                })
            })
            .await;

        assert!(response.is_ok());
        let stats = cache.stats();
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.hits, 0);

        // Second call (hit)
        let response2 = cache
            .get_or_compute("test query", async {
                // This should not be called
                panic!("Cache should have hit");
            })
            .await;

        assert!(response2.is_ok());
        let stats = cache.stats();
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.hits, 1);
        assert!(stats.hit_rate > 0.0);
    }

    #[test]
    fn test_cache_clear() {
        let cache = ResponseCache::new(100, Duration::from_secs(3600)).unwrap();
        cache.clear();

        let stats = cache.stats();
        assert_eq!(stats.size, 0);
    }

    #[test]
    fn test_query_hashing_consistency() {
        let query = "deployment procedure";
        let hash1 = hash_query(&normalize_query(query));
        let hash2 = hash_query(&normalize_query(query));

        assert_eq!(hash1, hash2, "Hash should be consistent");
    }
}
