//! Enhanced inference endpoints with comprehensive context caching

use axum::Json;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::OnceCell;

use crate::cache::{
    BeamSearchConfig, CacheConfig, CacheManager, CacheValue, WindowConfig,
    config::ContextCacheConfig,
    manager::PersistenceConfig,
    metrics::{AlertThresholds, CacheMetrics, CacheOperationType, OperationMetrics},
    speculative::TokenLogit,
};
use crate::{ApiError, ApiResult};

/// Enhanced inference request with caching parameters
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CachedInferenceRequest {
    /// Model ID to use for inference
    pub model: String,
    /// Input prompt
    pub prompt: String,
    /// Maximum tokens to generate
    pub max_tokens: Option<u32>,
    /// Temperature for sampling
    pub temperature: Option<f32>,
    /// Whether to stream the response
    pub stream: Option<bool>,
    /// Enable context caching
    pub enable_caching: Option<bool>,
    /// Enable speculative decoding
    pub enable_speculation: Option<bool>,
    /// Enable sliding window attention
    pub enable_sliding_window: Option<bool>,
    /// Cache key override
    pub cache_key: Option<String>,
    /// Request ID for tracking
    pub request_id: Option<String>,
}

/// Enhanced inference response with caching metadata
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CachedInferenceResponse {
    /// Generated text
    pub text: String,
    /// Finish reason
    pub finish_reason: String,
    /// Usage statistics
    pub usage: TokenUsage,
    /// Cache statistics for this request
    pub cache_stats: CacheRequestStats,
    /// Performance metrics
    pub performance: PerformanceMetrics,
    /// Request ID
    pub request_id: String,
}

/// Token usage statistics
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TokenUsage {
    /// Prompt tokens
    pub prompt_tokens: u32,
    /// Completion tokens
    pub completion_tokens: u32,
    /// Total tokens
    pub total_tokens: u32,
    /// Cached tokens (reused from cache)
    pub cached_tokens: u32,
}

/// Cache-related statistics for a request
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CacheRequestStats {
    /// Whether cache was used
    pub cache_used: bool,
    /// Cache hit rate for this request
    pub hit_rate: f64,
    /// Number of cache hits
    pub cache_hits: u32,
    /// Number of cache misses
    pub cache_misses: u32,
    /// Speculative decoding hits
    pub speculation_hits: u32,
    /// Radix tree prefix matches
    pub prefix_matches: u32,
    /// Sliding window cache reuse
    pub window_reuse: u32,
}

/// Performance metrics for a request
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PerformanceMetrics {
    /// Total inference time in milliseconds
    pub total_time_ms: f64,
    /// Time spent on caching operations in milliseconds
    pub cache_time_ms: f64,
    /// Time spent on actual inference in milliseconds
    pub inference_time_ms: f64,
    /// Tokens per second
    pub tokens_per_second: f64,
    /// Memory usage in bytes
    pub memory_usage_bytes: usize,
}

/// Global cache manager instance
static CACHE_MANAGER: OnceCell<Arc<CacheManager>> = OnceCell::const_new();
static CACHE_METRICS: OnceCell<Arc<CacheMetrics>> = OnceCell::const_new();

/// Initialize the cache system
pub async fn initialize_cache_system(config: ContextCacheConfig) -> ApiResult<()> {
    let cache_config = CacheConfig {
        l1_capacity: config.kv_cache.l1_capacity,
        l2_capacity: config.kv_cache.l2_capacity,
        l3_capacity_bytes: config.kv_cache.l3_capacity_bytes,
        max_memory_bytes: config.kv_cache.max_memory_bytes,
        eviction_policy: config.kv_cache.eviction_policy,
        ttl_seconds: config.kv_cache.ttl_seconds,
        enable_compression: config.kv_cache.enable_compression,
        cleanup_interval_seconds: config.kv_cache.cleanup_interval_seconds,
    };

    let beam_config = BeamSearchConfig {
        beam_width: config.speculative.beam_width,
        max_length: config.speculative.max_length,
        early_stopping_threshold: config.speculative.early_stopping_threshold,
        length_penalty: config.speculative.length_penalty,
        temperature: config.speculative.temperature,
        top_k: config.speculative.top_k,
        top_p: config.speculative.top_p,
        repetition_penalty: config.speculative.repetition_penalty,
    };

    let window_config = WindowConfig {
        window_size: config.sliding_window.window_size,
        window_overlap: config.sliding_window.window_overlap,
        max_cached_windows: config.sliding_window.max_cached_windows,
        use_flash_attention: config.sliding_window.use_flash_attention,
        attention_block_size: config.sliding_window.attention_block_size,
        gradient_checkpointing: config.sliding_window.gradient_checkpointing,
    };

    let persistence_config = PersistenceConfig {
        cache_dir: config.persistence.cache_dir,
        auto_persist: config.persistence.enabled,
        persist_interval_seconds: config.persistence.persist_interval_seconds,
        enable_compression: config.persistence.compression.enabled,
        max_file_size_bytes: config.persistence.max_file_size_bytes,
        backup_count: config.persistence.backup_count,
    };

    let cache_manager =
        CacheManager::new(cache_config, persistence_config, beam_config, window_config).await?;

    let alert_thresholds = AlertThresholds {
        min_hit_rate: config.monitoring.alerts.hit_rate_threshold,
        max_latency_us: config.monitoring.alerts.latency_threshold_us as f64,
        max_memory_bytes: config.monitoring.alerts.memory_threshold_percent as usize * 1024 * 1024, // Convert to bytes
        max_error_rate: config.monitoring.alerts.error_rate_threshold,
    };

    let cache_metrics = CacheMetrics::new(10000, alert_thresholds)
        .map_err(|e| ApiError::internal(format!("Failed to initialize metrics: {e}")))?;

    if config.monitoring.enabled {
        cache_metrics
            .start_background_collection(config.monitoring.collection_interval_seconds)
            .await;
    }

    CACHE_MANAGER
        .set(Arc::new(cache_manager))
        .map_err(|_| ApiError::internal("Cache manager already initialized"))?;

    CACHE_METRICS
        .set(Arc::new(cache_metrics))
        .map_err(|_| ApiError::internal("Cache metrics already initialized"))?;

    Ok(())
}

/// Get the cache manager instance
fn get_cache_manager() -> ApiResult<&'static Arc<CacheManager>> {
    CACHE_MANAGER
        .get()
        .ok_or_else(|| ApiError::internal("Cache system not initialized"))
}

/// Get the cache metrics instance
fn get_cache_metrics() -> ApiResult<&'static Arc<CacheMetrics>> {
    CACHE_METRICS
        .get()
        .ok_or_else(|| ApiError::internal("Cache metrics not initialized"))
}

/// Enhanced inference handler with context caching
pub async fn cached_inference(
    Json(request): Json<CachedInferenceRequest>,
) -> ApiResult<Json<CachedInferenceResponse>> {
    let start_time = Instant::now();
    let request_id = request
        .request_id
        .clone()
        .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());

    let cache_manager = get_cache_manager()?;
    let cache_metrics = get_cache_metrics()?;

    // Initialize request statistics
    let mut cache_stats = CacheRequestStats {
        cache_used: request.enable_caching.unwrap_or(true),
        hit_rate: 0.0,
        cache_hits: 0,
        cache_misses: 0,
        speculation_hits: 0,
        prefix_matches: 0,
        window_reuse: 0,
    };

    let cache_key = request
        .cache_key
        .clone()
        .unwrap_or_else(|| generate_cache_key(&request.model, &request.prompt));

    // Try cache lookup first if caching is enabled
    let mut cached_response = None;
    if cache_stats.cache_used {
        let cache_start = Instant::now();

        if let Some(cached_value) = cache_manager.get(&cache_key).await {
            cache_stats.cache_hits += 1;
            cached_response = Some(reconstruct_response_from_cache(&cached_value, &request_id));

            // Record cache hit metrics
            cache_metrics.record_operation(OperationMetrics {
                operation: CacheOperationType::Get,
                duration: cache_start.elapsed(),
                success: true,
                hit: Some(true),
                data_size: Some(cached_value.size_bytes),
                timestamp: cache_start,
                labels: std::collections::HashMap::new(),
            });
        } else {
            cache_stats.cache_misses += 1;

            // Record cache miss metrics
            cache_metrics.record_operation(OperationMetrics {
                operation: CacheOperationType::Get,
                duration: cache_start.elapsed(),
                success: true,
                hit: Some(false),
                data_size: None,
                timestamp: cache_start,
                labels: std::collections::HashMap::new(),
            });
        }
    }

    let inference_start = Instant::now();
    let mut response_text = String::new();
    let mut finish_reason = "stop".to_string();

    // If we have a cached response, use it
    let is_cached = cached_response.is_some();
    if let Some(cached_resp) = cached_response {
        response_text = cached_resp.text;
        finish_reason = cached_resp.finish_reason;
    } else {
        // Perform inference with optional enhancements
        let tokens = tokenize_prompt(&request.prompt)?;

        if request.enable_speculation.unwrap_or(false) {
            // Use speculative decoding
            let generated_tokens = cache_manager
                .speculative_decode(
                    tokens.clone(),
                    request.max_tokens.unwrap_or(100) as usize,
                    create_mock_logits_function(request.temperature.unwrap_or(1.0)),
                )
                .await;

            cache_stats.speculation_hits +=
                (generated_tokens.len().saturating_sub(tokens.len())) as u32;
            response_text = detokenize(&generated_tokens);
        } else if request.enable_sliding_window.unwrap_or(false) {
            // Use sliding window attention
            let sequence = tokens_to_sequence(&tokens);
            let processed = cache_manager
                .sliding_window_attention(&sequence, 64, 8)
                .await;

            cache_stats.window_reuse += estimate_window_reuse(&processed);
            response_text = sequence_to_text(&processed);
        } else {
            // Standard inference with potential prefix matching
            response_text = perform_standard_inference(&request, &mut cache_stats).await?;
        }

        // Cache the result if caching is enabled
        if cache_stats.cache_used {
            let cache_value = create_cache_value(&response_text, &tokens);
            if let Err(e) = cache_manager.put(cache_key, cache_value).await {
                tracing::warn!("Failed to cache inference result: {}", e);
            }
        }
    }

    let inference_time = inference_start.elapsed();
    let total_time = start_time.elapsed();

    // Calculate token counts
    let prompt_tokens = count_tokens(&request.prompt);
    let completion_tokens = count_tokens(&response_text);
    let cached_tokens = if is_cached { completion_tokens } else { 0 };

    let usage = TokenUsage {
        prompt_tokens,
        completion_tokens,
        total_tokens: prompt_tokens + completion_tokens,
        cached_tokens,
    };

    // Calculate hit rate
    cache_stats.hit_rate = if cache_stats.cache_hits + cache_stats.cache_misses > 0 {
        f64::from(cache_stats.cache_hits)
            / f64::from(cache_stats.cache_hits + cache_stats.cache_misses)
            * 100.0
    } else {
        0.0
    };

    // Performance metrics
    let cache_time = total_time.saturating_sub(inference_time);
    let performance = PerformanceMetrics {
        total_time_ms: total_time.as_secs_f64() * 1000.0,
        cache_time_ms: cache_time.as_secs_f64() * 1000.0,
        inference_time_ms: inference_time.as_secs_f64() * 1000.0,
        tokens_per_second: if total_time.as_secs_f64() > 0.0 {
            f64::from(completion_tokens) / total_time.as_secs_f64()
        } else {
            0.0
        },
        memory_usage_bytes: estimate_memory_usage(&response_text),
    };

    // Update cache manager statistics
    let manager_stats = cache_manager.stats().await;
    cache_metrics.update_cache_stats(&manager_stats);

    let response = CachedInferenceResponse {
        text: response_text,
        finish_reason,
        usage,
        cache_stats,
        performance,
        request_id,
    };

    Ok(Json(response))
}

/// Get cache statistics endpoint
pub async fn cache_stats() -> ApiResult<Json<serde_json::Value>> {
    let cache_manager = get_cache_manager()?;
    let cache_metrics = get_cache_metrics()?;

    let manager_stats = cache_manager.stats().await;
    let efficiency_score = cache_metrics.calculate_efficiency_score(&manager_stats);
    let alerts = cache_metrics.check_alerts(&manager_stats);
    let trends = cache_metrics.get_performance_trends();

    let stats = serde_json::json!({
        "cache_manager": {
            "kv_cache": {
                "hits": manager_stats.kv_cache.hits,
                "misses": manager_stats.kv_cache.misses,
                "hit_rate": manager_stats.kv_cache.hit_rate(),
                "memory_bytes": manager_stats.kv_cache.memory_bytes,
                "evictions": manager_stats.kv_cache.evictions,
                "l1_hits": manager_stats.kv_cache.l1_hits,
                "l2_hits": manager_stats.kv_cache.l2_hits,
                "l3_hits": manager_stats.kv_cache.l3_hits,
            },
            "radix_tree": {
                "prefix_matches": manager_stats.radix_tree.prefix_matches,
                "sharing_hits": manager_stats.radix_tree.sharing_hits,
                "memory_saved": manager_stats.radix_tree.memory_saved,
                "active_prefixes": manager_stats.radix_tree.active_prefixes,
                "avg_prefix_length": manager_stats.radix_tree.avg_prefix_length,
            },
            "speculation": {
                "hits": manager_stats.speculation.speculation_hits,
                "misses": manager_stats.speculation.speculation_misses,
                "hit_rate": manager_stats.speculation.hit_rate(),
                "tokens_generated": manager_stats.speculation.speculative_tokens,
                "avg_depth": manager_stats.speculation.avg_speculation_depth,
                "early_stops": manager_stats.speculation.early_stops,
            },
            "sliding_window": {
                "hits": manager_stats.sliding_window.window_hits,
                "misses": manager_stats.sliding_window.window_misses,
                "hit_rate": manager_stats.sliding_window.hit_rate(),
                "memory_bytes": manager_stats.sliding_window.total_memory_bytes,
                "evictions": manager_stats.sliding_window.windows_evicted,
                "utilization": manager_stats.sliding_window.avg_window_utilization,
            }
        },
        "overall": {
            "hit_rate": manager_stats.overall_hit_rate,
            "total_memory_bytes": manager_stats.total_memory_bytes,
            "operations_per_second": manager_stats.operations_per_second,
            "avg_response_time_us": manager_stats.avg_response_time_us,
            "efficiency_score": efficiency_score,
            "background_tasks_active": manager_stats.background_tasks_active,
        },
        "alerts": alerts.iter().map(|alert| {
            serde_json::json!({
                "severity": format!("{:?}", alert.severity),
                "message": alert.message,
                "metric": alert.metric,
                "current_value": alert.current_value,
                "threshold_value": alert.threshold_value,
            })
        }).collect::<Vec<_>>(),
        "trends": {
            "window_seconds": trends.window.as_secs(),
            "data_points": trends.hit_rate_trend.len(),
        }
    });

    Ok(Json(stats))
}

/// Clear all caches endpoint
pub async fn clear_cache() -> ApiResult<Json<serde_json::Value>> {
    let cache_manager = get_cache_manager()?;

    cache_manager.clear_all().await?;

    Ok(Json(serde_json::json!({
        "status": "success",
        "message": "All caches cleared successfully"
    })))
}

/// Warm cache endpoint
pub async fn warm_cache(
    Json(_request): Json<serde_json::Value>,
) -> ApiResult<Json<serde_json::Value>> {
    let cache_manager = get_cache_manager()?;

    let strategy = crate::cache::manager::WarmingStrategy::FrequentPatterns; // Default strategy
    let warmed_count = cache_manager.warm_cache(strategy).await?;

    Ok(Json(serde_json::json!({
        "status": "success",
        "warmed_entries": warmed_count,
        "message": format!("Cache warmed with {} entries", warmed_count)
    })))
}

// Helper functions

fn generate_cache_key(model: &str, prompt: &str) -> String {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    let mut hasher = DefaultHasher::new();
    model.hash(&mut hasher);
    prompt.hash(&mut hasher);
    format!("inference:{}:{:016x}", model, hasher.finish())
}

fn tokenize_prompt(prompt: &str) -> ApiResult<Vec<u32>> {
    // Simplified tokenization - in production would use actual tokenizer
    Ok(prompt
        .split_whitespace()
        .enumerate()
        .map(|(i, _)| i as u32)
        .collect())
}

fn detokenize(tokens: &[u32]) -> String {
    // Simplified detokenization
    format!("Generated text from {} tokens", tokens.len())
}

fn tokens_to_sequence(tokens: &[u32]) -> Vec<f32> {
    // Convert tokens to float sequence for attention computation
    tokens.iter().map(|&t| t as f32 * 0.1).collect()
}

fn sequence_to_text(sequence: &[f32]) -> String {
    format!(
        "Attention processed sequence with {} elements",
        sequence.len()
    )
}

fn count_tokens(text: &str) -> u32 {
    text.split_whitespace().count() as u32
}

fn create_cache_value(text: &str, tokens: &[u32]) -> CacheValue {
    let keys: Vec<f32> = tokens.iter().map(|&t| t as f32).collect();
    let values: Vec<f32> = text.chars().map(|c| c as u32 as f32).collect();

    CacheValue::new(keys, values, vec![tokens.len(), 1])
}

fn reconstruct_response_from_cache(
    cache_value: &CacheValue,
    request_id: &str,
) -> CachedInferenceResponse {
    let text = format!(
        "Cached response with {} keys and {} values",
        cache_value.keys.len(),
        cache_value.values.len()
    );

    CachedInferenceResponse {
        text,
        finish_reason: "cached".to_string(),
        usage: TokenUsage {
            prompt_tokens: 0,
            completion_tokens: cache_value.keys.len() as u32,
            total_tokens: cache_value.keys.len() as u32,
            cached_tokens: cache_value.keys.len() as u32,
        },
        cache_stats: CacheRequestStats {
            cache_used: true,
            hit_rate: 100.0,
            cache_hits: 1,
            cache_misses: 0,
            speculation_hits: 0,
            prefix_matches: 0,
            window_reuse: 0,
        },
        performance: PerformanceMetrics {
            total_time_ms: 0.1, // Very fast cache hit
            cache_time_ms: 0.1,
            inference_time_ms: 0.0,
            tokens_per_second: 10000.0, // Very fast
            memory_usage_bytes: cache_value.size_bytes,
        },
        request_id: request_id.to_string(),
    }
}

fn create_mock_logits_function(temperature: f32) -> impl Fn(&[u32]) -> Vec<TokenLogit> {
    move |tokens: &[u32]| {
        let base = tokens.last().unwrap_or(&0);
        vec![
            TokenLogit {
                token: base + 1,
                logit: 2.0 / temperature,
                probability: 0.6,
            },
            TokenLogit {
                token: base + 2,
                logit: 1.0 / temperature,
                probability: 0.3,
            },
            TokenLogit {
                token: base + 3,
                logit: 0.5 / temperature,
                probability: 0.1,
            },
        ]
    }
}

fn estimate_window_reuse(sequence: &[f32]) -> u32 {
    // Estimate how much of the sequence was reused from cache
    (sequence.len() / 4) as u32 // Assume 25% reuse
}

fn estimate_memory_usage(text: &str) -> usize {
    text.len() * 4 + std::mem::size_of::<String>() // Rough estimate
}

async fn perform_standard_inference(
    request: &CachedInferenceRequest,
    cache_stats: &mut CacheRequestStats,
) -> ApiResult<String> {
    // Mock inference - in production this would call actual model
    let response_text = format!(
        "Mock response to: '{}' with model '{}'",
        request.prompt, request.model
    );

    // Simulate some prefix matching
    if request.prompt.len() > 10 {
        cache_stats.prefix_matches += 1;
    }

    Ok(response_text)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cache::config::ContextCacheConfig;

    #[tokio::test]
    async fn test_cache_key_generation() {
        let key1 = generate_cache_key("model1", "hello world");
        let key2 = generate_cache_key("model1", "hello world");
        let key3 = generate_cache_key("model2", "hello world");

        assert_eq!(key1, key2); // Same inputs should generate same key
        assert_ne!(key1, key3); // Different inputs should generate different keys
        assert!(key1.starts_with("inference:model1:"));
    }

    #[test]
    fn test_tokenization() {
        let tokens = tokenize_prompt("hello world test").unwrap();
        assert_eq!(tokens.len(), 3);
        assert_eq!(tokens[0], 0);
        assert_eq!(tokens[1], 1);
        assert_eq!(tokens[2], 2);
    }

    #[test]
    fn test_token_counting() {
        assert_eq!(count_tokens("hello world"), 2);
        assert_eq!(count_tokens("single"), 1);
        assert_eq!(count_tokens(""), 0);
        assert_eq!(count_tokens("  multiple   spaces  "), 2);
    }

    #[test]
    fn test_cache_value_creation() {
        let text = "test response";
        let tokens = vec![1, 2, 3];
        let cache_value = create_cache_value(text, &tokens);

        assert_eq!(cache_value.keys.len(), 3);
        assert_eq!(cache_value.values.len(), text.len());
        assert_eq!(cache_value.shape, vec![3, 1]);
    }

    #[test]
    fn test_response_reconstruction() {
        let cache_value = CacheValue::new(vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0], vec![3, 1]);
        let response = reconstruct_response_from_cache(&cache_value, "test-request-123");

        assert_eq!(response.request_id, "test-request-123");
        assert_eq!(response.finish_reason, "cached");
        assert!(response.cache_stats.cache_used);
        assert_eq!(response.cache_stats.hit_rate, 100.0);
        assert_eq!(response.usage.cached_tokens, 3);
    }

    #[tokio::test]
    async fn test_cache_initialization() {
        let config = ContextCacheConfig::development();

        // Should not fail with development config
        // Note: This would require proper async test setup in production
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_memory_usage_estimation() {
        let text = "hello";
        let usage = estimate_memory_usage(text);

        // Should be greater than just the text length
        assert!(usage > text.len());
        assert!(usage > 0);
    }

    #[test]
    fn test_mock_logits_function() {
        let logits_fn = create_mock_logits_function(1.0);
        let tokens = vec![5, 10, 15];
        let logits = logits_fn(&tokens);

        assert_eq!(logits.len(), 3);
        assert_eq!(logits[0].token, 16); // last token (15) + 1
        assert!(logits[0].probability > logits[1].probability);
        assert!(logits[1].probability > logits[2].probability);

        // Test with temperature scaling
        let logits_fn_cold = create_mock_logits_function(0.5);
        let logits_cold = logits_fn_cold(&tokens);

        // Lower temperature should increase logits (making distribution sharper)
        assert!(logits_cold[0].logit > logits[0].logit);
    }

    #[test]
    fn test_window_reuse_estimation() {
        let sequence = vec![1.0; 100];
        let reuse = estimate_window_reuse(&sequence);

        assert_eq!(reuse, 25); // Should be 25% of sequence length

        let small_sequence = vec![1.0; 4];
        let small_reuse = estimate_window_reuse(&small_sequence);
        assert_eq!(small_reuse, 1);
    }
}