//! Query optimization for RAG system
//!
//! Provides intelligent query enhancement through:
//! - Intent detection (factual, comparison, how-to, etc)
//! - Query rewriting and normalization
//! - Synonym expansion for better matching
//! - Context injection from conversation history

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::error::Result;

/// Query intent classification
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum QueryIntent {
    /// Factual question (what, who, where, when)
    Factual,
    /// Comparison question (difference, comparison, vs)
    Comparison,
    /// How-to question (how, procedure, steps)
    HowTo,
    /// Why question (reason, cause, explanation)
    Why,
    /// Explanation request (explain, clarify, elaborate)
    Explanation,
    /// General query (unknown intent)
    General,
}

impl QueryIntent {
    /// Get human-readable name
    pub fn name(&self) -> &'static str {
        match self {
            QueryIntent::Factual => "Factual",
            QueryIntent::Comparison => "Comparison",
            QueryIntent::HowTo => "How-to",
            QueryIntent::Why => "Why",
            QueryIntent::Explanation => "Explanation",
            QueryIntent::General => "General",
        }
    }
}

/// Synonym mapping for domain-specific terms
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SynonymProvider {
    /// Map of terms to their synonyms
    synonyms: HashMap<String, Vec<String>>,
}

impl SynonymProvider {
    /// Create new synonym provider
    pub fn new() -> Self {
        let mut synonyms = HashMap::new();

        // Kubernetes synonyms
        synonyms.insert(
            "kubernetes".to_string(),
            vec![
                "k8s".to_string(),
                "k8".to_string(),
                "orchestrator".to_string(),
            ],
        );
        synonyms.insert(
            "pod".to_string(),
            vec!["container".to_string(), "application".to_string()],
        );
        synonyms.insert(
            "service".to_string(),
            vec!["endpoint".to_string(), "exposure".to_string()],
        );
        synonyms.insert(
            "deployment".to_string(),
            vec!["workload".to_string(), "application".to_string()],
        );

        // Generic synonyms
        synonyms.insert(
            "help".to_string(),
            vec!["assist".to_string(), "support".to_string()],
        );
        synonyms.insert(
            "show".to_string(),
            vec!["display".to_string(), "list".to_string()],
        );
        synonyms.insert(
            "create".to_string(),
            vec![
                "make".to_string(),
                "new".to_string(),
                "provision".to_string(),
            ],
        );
        synonyms.insert(
            "delete".to_string(),
            vec!["remove".to_string(), "destroy".to_string()],
        );

        Self { synonyms }
    }

    /// Add custom synonym
    pub fn add_synonym(&mut self, term: String, synonyms: Vec<String>) {
        self.synonyms.insert(term, synonyms);
    }

    /// Get synonyms for a term
    pub fn get_synonyms(&self, term: &str) -> Option<&Vec<String>> {
        self.synonyms.get(&term.to_lowercase())
    }

    /// Expand query with synonyms
    pub fn expand_query(&self, query: &str) -> String {
        let mut expanded = query.to_string();

        for (term, synonyms) in &self.synonyms {
            if query.to_lowercase().contains(term) {
                // Add synonyms as alternatives
                expanded.push_str(&format!(
                    " OR {} OR {}",
                    synonyms.first().unwrap_or(&String::new()),
                    synonyms.get(1).map(|s| s.as_str()).unwrap_or("")
                ));
            }
        }

        expanded
    }
}

impl Default for SynonymProvider {
    fn default() -> Self {
        Self::new()
    }
}

/// Intent detector for classifying queries
pub struct IntentDetector;

impl IntentDetector {
    /// Detect intent from query
    pub fn detect(query: &str) -> QueryIntent {
        let query_lower = query.to_lowercase();

        // How-to indicators
        let how_to_indicators = [
            "how do",
            "how to",
            "how can",
            "how would",
            "steps",
            "procedure",
            "process",
            "way to",
            "method",
            "guide",
            "tutorial",
        ];

        for indicator in &how_to_indicators {
            if query_lower.contains(indicator) {
                return QueryIntent::HowTo;
            }
        }

        // Why indicators
        let why_indicators = ["why", "reason for", "cause", "explanation for"];

        for indicator in &why_indicators {
            if query_lower.contains(indicator) {
                return QueryIntent::Why;
            }
        }

        // Comparison indicators
        let comparison_indicators = [
            "difference",
            "compare",
            "vs",
            "versus",
            "different from",
            "similar to",
            "like",
            "compared to",
            "contrast",
        ];

        for indicator in &comparison_indicators {
            if query_lower.contains(indicator) {
                return QueryIntent::Comparison;
            }
        }

        // Explanation indicators
        let explanation_indicators = ["explain", "clarify", "elaborate", "describe", "tell me"];

        for indicator in &explanation_indicators {
            if query_lower.contains(indicator) {
                return QueryIntent::Explanation;
            }
        }

        // Factual indicators
        let factual_indicators = ["what", "who", "where", "when", "which"];

        for indicator in &factual_indicators {
            if query_lower.starts_with(indicator) {
                return QueryIntent::Factual;
            }
        }

        QueryIntent::General
    }
}

/// Optimized query with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizedQuery {
    /// Original query
    pub original: String,
    /// Optimized/rewritten query
    pub optimized: String,
    /// Detected intent
    pub intent: QueryIntent,
    /// Confidence in intent detection (0.0-1.0)
    pub intent_confidence: f32,
    /// Extracted key terms
    pub key_terms: Vec<String>,
    /// Added context
    pub context_injected: bool,
}

impl OptimizedQuery {
    /// Create new optimized query
    pub fn new(original: String, optimized: String, intent: QueryIntent) -> Self {
        Self {
            original,
            optimized,
            intent,
            intent_confidence: 0.8,
            key_terms: Vec::new(),
            context_injected: false,
        }
    }

    /// Set intent confidence
    pub fn with_confidence(mut self, confidence: f32) -> Self {
        self.intent_confidence = confidence.clamp(0.0, 1.0);
        self
    }

    /// Set key terms
    pub fn with_key_terms(mut self, terms: Vec<String>) -> Self {
        self.key_terms = terms;
        self
    }

    /// Mark context as injected
    pub fn mark_context_injected(mut self) -> Self {
        self.context_injected = true;
        self
    }
}

/// Query optimizer engine
pub struct QueryOptimizer {
    synonym_provider: SynonymProvider,
}

impl QueryOptimizer {
    /// Create new query optimizer
    pub fn new() -> Self {
        Self {
            synonym_provider: SynonymProvider::new(),
        }
    }

    /// Create with custom synonym provider
    pub fn with_synonyms(synonym_provider: SynonymProvider) -> Self {
        Self { synonym_provider }
    }

    /// Extract key terms from query
    pub fn extract_key_terms(&self, query: &str) -> Vec<String> {
        // Remove common words
        let stopwords = [
            "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with",
            "by", "from", "up", "about", "into", "through", "during", "is", "are", "was", "were",
            "be", "been", "being", "do", "does", "did", "how", "what", "when", "where", "why",
            "which", "who", "my", "your",
        ];

        query
            .to_lowercase()
            .split_whitespace()
            .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()))
            .filter(|w| !w.is_empty() && !stopwords.contains(w))
            .map(String::from)
            .collect::<Vec<_>>()
            .into_iter()
            .collect::<std::collections::HashSet<_>>()
            .into_iter()
            .collect()
    }

    /// Normalize query (lowercase, trim, remove extra spaces)
    pub fn normalize_query(&self, query: &str) -> String {
        query
            .to_lowercase()
            .split_whitespace()
            .collect::<Vec<_>>()
            .join(" ")
            .trim()
            .to_string()
    }

    /// Rewrite query with improvements
    pub fn rewrite_query(&self, query: &str) -> String {
        let normalized = self.normalize_query(query);

        // Add synonym expansion
        self.synonym_provider.expand_query(&normalized)
    }

    /// Optimize query with full analysis
    pub fn optimize(&self, query: &str) -> Result<OptimizedQuery> {
        let original = query.to_string();
        let normalized = self.normalize_query(query);
        let rewritten = self.rewrite_query(query);
        let intent = IntentDetector::detect(&normalized);
        let key_terms = self.extract_key_terms(&normalized);

        // Adjust confidence based on intent certainty
        let confidence = match intent {
            QueryIntent::HowTo | QueryIntent::Why | QueryIntent::Factual => 0.95,
            QueryIntent::Explanation | QueryIntent::Comparison => 0.85,
            QueryIntent::General => 0.6,
        };

        Ok(OptimizedQuery::new(original, rewritten, intent)
            .with_confidence(confidence)
            .with_key_terms(key_terms))
    }

    /// Inject context from conversation history
    pub fn inject_context(&self, query: &str, context: &str) -> String {
        format!("{}\n\nContext: {}", query, context)
    }

    /// Full optimization pipeline
    pub fn optimize_with_context(
        &self,
        query: &str,
        context: Option<&str>,
    ) -> Result<OptimizedQuery> {
        let mut optimized = self.optimize(query)?;

        if let Some(ctx) = context {
            let enriched = self.inject_context(&optimized.optimized, ctx);
            optimized.optimized = enriched;
            optimized = optimized.mark_context_injected();
        }

        Ok(optimized)
    }
}

impl Default for QueryOptimizer {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_query_intent_names() {
        assert_eq!(QueryIntent::Factual.name(), "Factual");
        assert_eq!(QueryIntent::HowTo.name(), "How-to");
        assert_eq!(QueryIntent::Why.name(), "Why");
    }

    #[test]
    fn test_synonym_provider_creation() {
        let provider = SynonymProvider::new();
        assert!(provider.get_synonyms("kubernetes").is_some());
    }

    #[test]
    fn test_synonym_lookup() {
        let provider = SynonymProvider::new();
        let synonyms = provider.get_synonyms("kubernetes");
        assert!(synonyms.is_some());
        assert!(synonyms.unwrap().contains(&"k8s".to_string()));
    }

    #[test]
    fn test_add_custom_synonym() {
        let mut provider = SynonymProvider::new();
        provider.add_synonym("custom".to_string(), vec!["alias".to_string()]);
        assert!(provider.get_synonyms("custom").is_some());
    }

    #[test]
    fn test_intent_detection_howto() {
        let intent = IntentDetector::detect("How do I deploy Kubernetes?");
        assert_eq!(intent, QueryIntent::HowTo);
    }

    #[test]
    fn test_intent_detection_why() {
        let intent = IntentDetector::detect("Why is scaling important?");
        assert_eq!(intent, QueryIntent::Why);
    }

    #[test]
    fn test_intent_detection_comparison() {
        let intent =
            IntentDetector::detect("What is the difference between Docker and Kubernetes?");
        assert_eq!(intent, QueryIntent::Comparison);
    }

    #[test]
    fn test_intent_detection_explanation() {
        let intent = IntentDetector::detect("Can you explain how services work?");
        assert_eq!(intent, QueryIntent::Explanation);
    }

    #[test]
    fn test_intent_detection_factual() {
        let intent = IntentDetector::detect("What is Kubernetes?");
        assert_eq!(intent, QueryIntent::Factual);
    }

    #[test]
    fn test_query_normalization() {
        let optimizer = QueryOptimizer::new();
        let normalized = optimizer.normalize_query("  What   is   KUBERNETES?  ");
        assert_eq!(normalized, "what is kubernetes?");
    }

    #[test]
    fn test_extract_key_terms() {
        let optimizer = QueryOptimizer::new();
        let terms = optimizer.extract_key_terms("How do I deploy Kubernetes applications?");
        assert!(terms.contains(&"deploy".to_string()));
        assert!(terms.contains(&"kubernetes".to_string()));
        assert!(!terms.contains(&"how".to_string())); // stopword
    }

    #[test]
    fn test_query_optimization() {
        let optimizer = QueryOptimizer::new();
        let optimized = optimizer.optimize("What is kubernetes?").unwrap();
        assert_eq!(optimized.intent, QueryIntent::Factual);
        assert!(!optimized.key_terms.is_empty());
        assert!(!optimized.context_injected);
    }

    #[test]
    fn test_optimize_with_confidence() {
        let optimizer = QueryOptimizer::new();
        let optimized = optimizer.optimize("How do I create a pod?").unwrap();
        assert!(optimized.intent_confidence > 0.9);
        assert_eq!(optimized.intent, QueryIntent::HowTo);
    }

    #[test]
    fn test_context_injection() {
        let optimizer = QueryOptimizer::new();
        let context = "We discussed Kubernetes basics earlier";
        let injected = optimizer.inject_context("Tell me more about services", context);
        assert!(injected.contains("Context:"));
        assert!(injected.contains("services"));
    }

    #[test]
    fn test_optimize_with_context() {
        let optimizer = QueryOptimizer::new();
        let context = "Previously discussed: orchestration platforms";
        let optimized = optimizer
            .optimize_with_context("Tell me more", Some(context))
            .unwrap();
        assert!(optimized.context_injected);
        assert!(optimized.optimized.contains("Context:"));
    }

    #[test]
    fn test_query_rewriting() {
        let optimizer = QueryOptimizer::new();
        let rewritten = optimizer.rewrite_query("What is kubernetes?");
        // Should contain expanded synonyms or normalized form
        assert!(!rewritten.is_empty());
    }

    #[test]
    fn test_optimized_query_building() {
        let opt_query =
            OptimizedQuery::new("original".into(), "optimized".into(), QueryIntent::Factual)
                .with_confidence(0.95)
                .with_key_terms(vec!["term1".into(), "term2".into()])
                .mark_context_injected();

        assert_eq!(opt_query.intent_confidence, 0.95);
        assert_eq!(opt_query.key_terms.len(), 2);
        assert!(opt_query.context_injected);
    }

    #[test]
    fn test_intent_confidence_clamping() {
        let opt_query =
            OptimizedQuery::new("original".into(), "optimized".into(), QueryIntent::General)
                .with_confidence(2.5); // Out of range

        assert_eq!(opt_query.intent_confidence, 1.0);
    }

    #[test]
    fn test_general_intent_fallback() {
        let intent = IntentDetector::detect("Random text without clear intent markers");
        assert_eq!(intent, QueryIntent::General);
    }

    #[test]
    fn test_case_insensitive_intent_detection() {
        let intent1 = IntentDetector::detect("HOW DO I DEPLOY?");
        let intent2 = IntentDetector::detect("how do i deploy?");
        assert_eq!(intent1, intent2);
    }

    #[test]
    fn test_synonym_expansion_with_multiple_terms() {
        let mut provider = SynonymProvider::new();
        provider.add_synonym(
            "container".to_string(),
            vec!["image".to_string(), "app".to_string()],
        );

        let expanded = provider.expand_query("kubernetes and container");
        assert!(expanded.contains("container"));
    }
}
