use std::sync::Arc;

use anyhow::Result;
use tracing::{debug, warn};

use crate::types::{LlmProvider, QueryType};

/// Classifies a natural-language search query into a `QueryType` so the
/// orchestrator can choose an appropriate search strategy.
pub struct QueryClassifier {
    llm: Arc<dyn LlmProvider>,
}

impl QueryClassifier {
    pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
        Self { llm }
    }

    /// Classify a query string using the LLM. Falls back to heuristics if
    /// the LLM response cannot be parsed.
    pub async fn classify(&self, query: &str) -> Result<QueryType> {
        // Try fast heuristic classification first to avoid unnecessary LLM calls.
        if let Some(qt) = self.heuristic_classify(query) {
            debug!(query = %query, query_type = ?qt, "Classified via heuristics");
            return Ok(qt);
        }

        // If LLM is unavailable, default to Semantic.
        if !self.llm.is_available() {
            debug!(query = %query, "LLM unavailable, defaulting to Semantic");
            return Ok(QueryType::Semantic);
        }

        // Fall back to LLM-based classification.
        let system_prompt = r#"You are a search query classifier for a code intelligence system.
Given a user query about a codebase, classify it into exactly ONE of these categories:

- ExactMatch: Looking for a specific symbol, function name, variable, or exact string.
  Examples: "find function handle_request", "where is TIMEOUT_MS defined"
- Semantic: Conceptual/meaning-based search, understanding intent rather than exact text.
  Examples: "how does authentication work", "find error handling logic"
- Pattern: Looking for code patterns, conventions, or structural idioms.
  Examples: "find all singleton patterns", "show builder pattern usage"
- Architecture: Questions about system structure, module relationships, data flow.
  Examples: "how do modules depend on each other", "what is the request lifecycle"
- Dependency: Questions about imports, dependencies, or what uses/is used by what.
  Examples: "what imports this module", "find all callers of parse()"

Respond with ONLY the category name, nothing else."#;

        let response = tokio::time::timeout(
            std::time::Duration::from_secs(30),
            self.llm.complete_with_system(system_prompt, query),
        )
        .await
        .map_err(|_| anyhow::anyhow!("LLM query classification timed out after 30s"))??;
        let trimmed = response.trim();

        let query_type = Self::parse_query_type(trimmed).unwrap_or_else(|| {
            warn!(
                query = %query,
                llm_response = %trimmed,
                "LLM returned unparseable query type, falling back to Semantic"
            );
            QueryType::Semantic
        });

        debug!(query = %query, query_type = ?query_type, "Classified via LLM");
        Ok(query_type)
    }

    /// Attempt a fast, no-LLM classification based on keyword patterns.
    fn heuristic_classify(&self, query: &str) -> Option<QueryType> {
        let lower = query.to_lowercase();

        // Dependency indicators
        if lower.contains("import")
            || lower.contains("depend")
            || lower.contains("caller")
            || lower.contains("callee")
            || lower.contains("used by")
            || lower.contains("uses ")
        {
            return Some(QueryType::Dependency);
        }

        // Architecture indicators
        if lower.contains("architect")
            || lower.contains("module")
            || lower.contains("data flow")
            || lower.contains("lifecycle")
            || lower.contains("how do .* relate")
        {
            return Some(QueryType::Architecture);
        }

        // Pattern indicators
        if lower.contains("pattern")
            || lower.contains("convention")
            || lower.contains("idiom")
            || lower.contains("all instances of")
        {
            return Some(QueryType::Pattern);
        }

        // Exact match indicators: quoted strings or very short queries that look like symbols
        if lower.contains('"')
            || lower.contains('\'')
            || lower.starts_with("find function ")
            || lower.starts_with("find fn ")
            || lower.starts_with("where is ")
            || lower.starts_with("definition of ")
            || lower.starts_with("def ")
        {
            return Some(QueryType::ExactMatch);
        }

        // If the query looks like a single symbol (no spaces, or camelCase/snake_case),
        // treat it as ExactMatch.
        let words: Vec<&str> = query.split_whitespace().collect();
        if words.len() == 1 && (query.contains('_') || query.chars().any(|c| c.is_uppercase())) {
            return Some(QueryType::ExactMatch);
        }

        // Cannot determine heuristically; defer to LLM.
        None
    }

    /// Parse a QueryType from a string (case-insensitive).
    fn parse_query_type(s: &str) -> Option<QueryType> {
        match s.to_lowercase().trim() {
            "exactmatch" | "exact_match" | "exact match" => Some(QueryType::ExactMatch),
            "semantic" => Some(QueryType::Semantic),
            "pattern" => Some(QueryType::Pattern),
            "architecture" => Some(QueryType::Architecture),
            "dependency" => Some(QueryType::Dependency),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::QueryType;
    use std::sync::Arc;

    struct MockLlm {
        response: String,
    }

    impl MockLlm {
        fn new(response: &str) -> Arc<Self> {
            Arc::new(Self { response: response.to_string() })
        }
    }

    #[async_trait::async_trait]
    impl LlmProvider for MockLlm {
        async fn complete(&self, _prompt: &str) -> anyhow::Result<String> {
            Ok(self.response.clone())
        }
        async fn complete_with_system(&self, _system: &str, _prompt: &str) -> anyhow::Result<String> {
            Ok(self.response.clone())
        }
    }

    // -- heuristic_classify tests --

    #[tokio::test]
    async fn test_heuristic_dependency_import() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("what imports this module"), Some(QueryType::Dependency));
    }

    #[tokio::test]
    async fn test_heuristic_dependency_caller() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("find all callers of parse"), Some(QueryType::Dependency));
    }

    #[tokio::test]
    async fn test_heuristic_architecture_module() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("what is the module lifecycle"), Some(QueryType::Architecture));
    }

    #[tokio::test]
    async fn test_heuristic_architecture_data_flow() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("describe the data flow"), Some(QueryType::Architecture));
    }

    #[tokio::test]
    async fn test_heuristic_pattern() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("find all singleton patterns"), Some(QueryType::Pattern));
    }

    #[tokio::test]
    async fn test_heuristic_exact_match_quotes() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify(r#"find "handle_request""#), Some(QueryType::ExactMatch));
    }

    #[tokio::test]
    async fn test_heuristic_exact_match_find_function() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("find function handle_request"), Some(QueryType::ExactMatch));
    }

    #[tokio::test]
    async fn test_heuristic_exact_match_snake_case_symbol() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("find_this_symbol"), Some(QueryType::ExactMatch));
    }

    #[tokio::test]
    async fn test_heuristic_exact_match_camel_case_symbol() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("HandleRequest"), Some(QueryType::ExactMatch));
    }

    #[tokio::test]
    async fn test_heuristic_returns_none_for_semantic() {
        let c = QueryClassifier::new(MockLlm::new(""));
        assert_eq!(c.heuristic_classify("how does authentication work in this codebase"), None);
    }

    // -- LLM fallback tests --

    #[tokio::test]
    async fn test_llm_classify_semantic() {
        let c = QueryClassifier::new(MockLlm::new("Semantic"));
        let qt = c.classify("how does authentication work").await.unwrap();
        assert_eq!(qt, QueryType::Semantic);
    }

    #[tokio::test]
    async fn test_llm_classify_exact_match_case_insensitive() {
        let c = QueryClassifier::new(MockLlm::new("EXACTMATCH"));
        let qt = c.classify("how does authentication work").await.unwrap();
        assert_eq!(qt, QueryType::ExactMatch);
    }

    #[tokio::test]
    async fn test_llm_classify_unparseable_falls_back_to_semantic() {
        let c = QueryClassifier::new(MockLlm::new("I'm not sure what category this is"));
        let qt = c.classify("how does authentication work").await.unwrap();
        assert_eq!(qt, QueryType::Semantic);
    }

    // -- parse_query_type tests --

    #[test]
    fn test_parse_query_type_all_variants() {
        assert_eq!(QueryClassifier::parse_query_type("exactmatch"), Some(QueryType::ExactMatch));
        assert_eq!(QueryClassifier::parse_query_type("exact_match"), Some(QueryType::ExactMatch));
        assert_eq!(QueryClassifier::parse_query_type("exact match"), Some(QueryType::ExactMatch));
        assert_eq!(QueryClassifier::parse_query_type("semantic"), Some(QueryType::Semantic));
        assert_eq!(QueryClassifier::parse_query_type("pattern"), Some(QueryType::Pattern));
        assert_eq!(QueryClassifier::parse_query_type("architecture"), Some(QueryType::Architecture));
        assert_eq!(QueryClassifier::parse_query_type("dependency"), Some(QueryType::Dependency));
        assert_eq!(QueryClassifier::parse_query_type("SEMANTIC"), Some(QueryType::Semantic));
        assert_eq!(QueryClassifier::parse_query_type("unknown"), None);
    }
}
