use std::sync::Arc;

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

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

/// Uses an LLM to rank search results by relevance to a query.
///
/// The ranker sends the query and a summary of the results to the LLM, asking
/// it to assign a relevance score (0.0 - 1.0) to each result. The results are
/// then sorted by descending score.
pub struct ResultRanker {
    llm: Arc<dyn LlmProvider>,
}

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

    /// Rank the given results by relevance to `query` using the LLM.
    ///
    /// If the LLM call fails or returns unparseable output, the results are
    /// returned in their original order with a warning logged.
    #[instrument(skip(self, results), fields(result_count = results.len()))]
    pub async fn rank(&self, query: &str, mut results: Vec<SearchResult>) -> Result<Vec<SearchResult>> {
        if results.is_empty() {
            return Ok(results);
        }

        // For very small result sets, skip the LLM call and just return as-is.
        if results.len() <= 2 {
            debug!("skipping LLM ranking for <= 2 results");
            return Ok(results);
        }

        // Skip LLM ranking if provider is unavailable.
        if !self.llm.is_available() {
            debug!("ResultRanker: LLM unavailable, returning results in original order");
            return Ok(results);
        }

        // Cap the number of results we send to the LLM to avoid token
        // overflow.
        let max_to_rank = 50;
        let truncated = results.len() > max_to_rank;
        let to_rank: Vec<_> = results.iter().take(max_to_rank).collect();

        let prompt = self.build_ranking_prompt(query, &to_rank);

        let system = "You are a search result ranker. Given a query and a list of code search \
                       results, assign a relevance score between 0.0 and 1.0 to each result. \
                       Respond ONLY with a JSON array of objects like: \
                       [{\"index\": 0, \"score\": 0.95}, {\"index\": 1, \"score\": 0.3}]. \
                       Order does not matter. Be precise.";

        let llm_result = tokio::time::timeout(
            std::time::Duration::from_secs(30),
            self.llm.complete_with_system(system, &prompt),
        )
        .await;
        match llm_result {
            Ok(Ok(response)) => {
                match self.parse_ranking_response(&response, results.len()) {
                    Ok(scores) => {
                        self.apply_scores(&mut results, &scores);
                        results.sort_by(|a, b| {
                            b.relevance_score
                                .partial_cmp(&a.relevance_score)
                                .unwrap_or(std::cmp::Ordering::Equal)
                        });
                        debug!("LLM ranking applied successfully");
                    }
                    Err(e) => {
                        warn!(error = %e, "failed to parse LLM ranking response; using original order");
                    }
                }
            }
            Ok(Err(e)) => {
                warn!(error = %e, "LLM ranking call failed; using original order");
            }
            Err(_) => {
                warn!("LLM ranking timed out after 30s; using original order");
            }
        }

        if truncated {
            debug!(
                total = results.len(),
                ranked = max_to_rank,
                "only top {} results were ranked by LLM",
                max_to_rank
            );
        }

        Ok(results)
    }

    // ---- internal ----

    /// Build the prompt that lists all results with their indices for the LLM.
    fn build_ranking_prompt(&self, query: &str, results: &[&SearchResult]) -> String {
        let mut prompt = format!(
            "Query: \"{}\"\n\nSearch Results:\n",
            query
        );

        for (i, result) in results.iter().enumerate() {
            let line_info = result
                .line_number
                .map(|n| format!(":{n}"))
                .unwrap_or_default();

            // Truncate content to keep the prompt manageable.
            let content = if result.content.len() > 200 {
                format!("{}...", &result.content[..200])
            } else {
                result.content.clone()
            };

            prompt.push_str(&format!(
                "[{}] {}{} | {}\n",
                i, result.file_path, line_info, content
            ));
        }

        prompt.push_str(
            "\nReturn a JSON array with a relevance score (0.0-1.0) for each result index.",
        );

        prompt
    }

    /// Parse the LLM's ranking response into a mapping from result index to
    /// score.
    fn parse_ranking_response(
        &self,
        response: &str,
        result_count: usize,
    ) -> Result<Vec<(usize, f64)>> {
        // Extract JSON array from response.
        let json_str = Self::extract_json_array(response)
            .context("no JSON array found in ranking response")?;

        let entries: Vec<serde_json::Value> = serde_json::from_str(&json_str)
            .context("failed to parse ranking JSON array")?;

        let mut scores = Vec::new();

        for entry in entries {
            let index = entry
                .get("index")
                .and_then(|v| v.as_u64())
                .map(|v| v as usize);
            let score = entry
                .get("score")
                .and_then(|v| v.as_f64());

            if let (Some(idx), Some(s)) = (index, score) {
                if idx < result_count {
                    // Clamp score to [0.0, 1.0].
                    let clamped = s.clamp(0.0, 1.0);
                    scores.push((idx, clamped));
                }
            }
        }

        if scores.is_empty() {
            anyhow::bail!("no valid score entries parsed from LLM response");
        }

        Ok(scores)
    }

    /// Apply the parsed scores to the results vector.
    fn apply_scores(&self, results: &mut [SearchResult], scores: &[(usize, f64)]) {
        for (idx, score) in scores {
            if *idx < results.len() {
                results[*idx].relevance_score = *score;
            }
        }
    }

    /// Find the first JSON array `[...]` in a potentially noisy string.
    fn extract_json_array(text: &str) -> Option<String> {
        let start = text.find('[')?;
        let end = text.rfind(']')?;
        if end > start {
            Some(text[start..=end].to_string())
        } else {
            None
        }
    }
}

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

    #[test]
    fn test_parse_ranking_response() {
        let ranker = ResultRanker {
            llm: Arc::new(MockLlm),
        };
        let response = r#"[{"index": 0, "score": 0.9}, {"index": 1, "score": 0.3}]"#;
        let scores = ranker.parse_ranking_response(response, 5).unwrap();
        assert_eq!(scores.len(), 2);
        assert_eq!(scores[0], (0, 0.9));
        assert_eq!(scores[1], (1, 0.3));
    }

    #[test]
    fn test_parse_ranking_clamps_scores() {
        let ranker = ResultRanker {
            llm: Arc::new(MockLlm),
        };
        let response = r#"[{"index": 0, "score": 1.5}, {"index": 1, "score": -0.2}]"#;
        let scores = ranker.parse_ranking_response(response, 5).unwrap();
        assert!((scores[0].1 - 1.0).abs() < f64::EPSILON);
        assert!((scores[1].1 - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_parse_ranking_ignores_out_of_range() {
        let ranker = ResultRanker {
            llm: Arc::new(MockLlm),
        };
        let response = r#"[{"index": 99, "score": 0.5}]"#;
        let result = ranker.parse_ranking_response(response, 5);
        // Out-of-range indices are filtered out, leaving no valid scores,
        // which causes the parser to return an error.
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_json_array() {
        let text = "Here are the scores: [{\"index\": 0, \"score\": 0.8}] done.";
        let result = ResultRanker::extract_json_array(text);
        assert!(result.is_some());
        assert!(result.unwrap().starts_with('['));
    }

    // Minimal mock for unit tests.
    struct MockLlm;

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