// LLM Dispatch - Distributed Reasoning
// Sends chunks to LLM providers for analysis and aggregates results

use std::sync::Arc;
use std::time::Instant;

use tracing::{debug, info};
use vapora_llm_router::providers::LLMClient;

use crate::metrics::DISPATCH_DURATION;
use crate::search::hybrid::ScoredChunk;
use crate::RLMError;

/// Dispatch configuration
#[derive(Debug, Clone)]
pub struct DispatchConfig {
    /// Include chunk content in prompt
    pub include_content: bool,
    /// Include chunk metadata in prompt
    pub include_metadata: bool,
    /// Maximum chunks per dispatch
    pub max_chunks_per_dispatch: usize,
    /// Aggregation strategy
    pub aggregation: AggregationStrategy,
}

impl Default for DispatchConfig {
    fn default() -> Self {
        Self {
            include_content: true,
            include_metadata: false,
            max_chunks_per_dispatch: 10,
            aggregation: AggregationStrategy::Concatenate,
        }
    }
}

/// Strategy for aggregating results from multiple LLM calls
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggregationStrategy {
    /// Concatenate all responses
    Concatenate,
    /// Take first response only
    FirstOnly,
    /// Use majority voting (for classification tasks)
    MajorityVote,
}

/// Dispatch result from a single LLM call
#[derive(Debug, Clone)]
pub struct DispatchResult {
    /// Response text from LLM
    pub text: String,
    /// Input tokens used
    pub input_tokens: u64,
    /// Output tokens generated
    pub output_tokens: u64,
    /// Finish reason
    pub finish_reason: String,
    /// Duration in milliseconds
    pub duration_ms: u64,
}

/// Aggregated dispatch results
#[derive(Debug, Clone)]
pub struct AggregatedResult {
    /// Combined response text
    pub text: String,
    /// Total input tokens across all calls
    pub total_input_tokens: u64,
    /// Total output tokens across all calls
    pub total_output_tokens: u64,
    /// Number of LLM calls made
    pub num_calls: usize,
    /// Total duration in milliseconds
    pub total_duration_ms: u64,
}

/// LLM dispatcher for distributed reasoning
pub struct LLMDispatcher {
    llm_client: Option<Arc<dyn LLMClient>>,
    config: DispatchConfig,
}

impl LLMDispatcher {
    /// Create a new dispatcher with an LLM client
    pub fn new(llm_client: Option<Arc<dyn LLMClient>>) -> Self {
        Self {
            llm_client,
            config: DispatchConfig::default(),
        }
    }

    /// Create with custom configuration
    pub fn with_config(llm_client: Option<Arc<dyn LLMClient>>, config: DispatchConfig) -> Self {
        Self { llm_client, config }
    }

    /// Dispatch chunks to LLM for analysis
    ///
    /// # Arguments
    /// - `query`: User query/task description
    /// - `chunks`: Relevant chunks from hybrid search
    ///
    /// # Returns
    /// Aggregated result from all LLM calls
    pub async fn dispatch(
        &self,
        query: &str,
        chunks: &[ScoredChunk],
    ) -> crate::Result<AggregatedResult> {
        let start = Instant::now();

        if chunks.is_empty() {
            return Ok(AggregatedResult {
                text: String::new(),
                total_input_tokens: 0,
                total_output_tokens: 0,
                num_calls: 0,
                total_duration_ms: 0,
            });
        }

        // Get LLM client
        let client = self
            .llm_client
            .as_ref()
            .ok_or_else(|| RLMError::DispatchError("LLM client not configured".to_string()))?;

        info!(
            "Dispatching {} chunks to LLM: provider={}",
            chunks.len(),
            client.provider_name()
        );

        // Split chunks into batches if needed
        let batches = self.split_into_batches(chunks);

        // Dispatch each batch
        let mut results = Vec::new();
        for (batch_idx, batch) in batches.iter().enumerate() {
            debug!(
                "Dispatching batch {}/{} with {} chunks",
                batch_idx + 1,
                batches.len(),
                batch.len()
            );

            let result = self.dispatch_batch(client.as_ref(), query, batch).await?;
            results.push(result);
        }

        // Aggregate results
        let aggregated = self.aggregate_results(results);

        let duration = start.elapsed();
        DISPATCH_DURATION
            .with_label_values(&[&client.provider_name()])
            .observe(duration.as_secs_f64());

        info!(
            "Dispatch completed: {} calls, {} input tokens, {} output tokens, {:?}",
            aggregated.num_calls,
            aggregated.total_input_tokens,
            aggregated.total_output_tokens,
            duration
        );

        Ok(aggregated)
    }

    /// Dispatch a single batch of chunks
    async fn dispatch_batch(
        &self,
        client: &dyn LLMClient,
        query: &str,
        chunks: &[&ScoredChunk],
    ) -> crate::Result<DispatchResult> {
        let start = Instant::now();

        // Build prompt with chunks
        let prompt = self.build_prompt(query, chunks);

        // Call LLM
        let response = client
            .complete(prompt, None)
            .await
            .map_err(|e| RLMError::DispatchError(format!("LLM call failed: {}", e)))?;

        let duration_ms = start.elapsed().as_millis() as u64;

        Ok(DispatchResult {
            text: response.text,
            input_tokens: response.input_tokens,
            output_tokens: response.output_tokens,
            finish_reason: response.finish_reason,
            duration_ms,
        })
    }

    /// Build prompt from query and chunks
    fn build_prompt(&self, query: &str, chunks: &[&ScoredChunk]) -> String {
        let mut prompt = format!("Query: {}\n\nRelevant information:\n\n", query);

        for (idx, chunk) in chunks.iter().enumerate() {
            prompt.push_str(&format!("=== Chunk {} ===\n", idx + 1));

            if self.config.include_content {
                prompt.push_str(&chunk.chunk.content);
                prompt.push_str("\n\n");
            }

            if self.config.include_metadata {
                prompt.push_str(&format!(
                    "Metadata: chunk_id={}, doc_id={}, BM25 score={:?}, semantic score={:?}\n\n",
                    chunk.chunk.chunk_id,
                    chunk.chunk.doc_id,
                    chunk.bm25_score,
                    chunk.semantic_score
                ));
            }
        }

        prompt.push_str("Based on the above information, please answer the query.\n");
        prompt
    }

    /// Split chunks into batches based on max_chunks_per_dispatch
    fn split_into_batches<'a>(&self, chunks: &'a [ScoredChunk]) -> Vec<Vec<&'a ScoredChunk>> {
        let mut batches = Vec::new();
        let mut current_batch = Vec::new();

        for chunk in chunks {
            current_batch.push(chunk);
            if current_batch.len() >= self.config.max_chunks_per_dispatch {
                batches.push(current_batch);
                current_batch = Vec::new();
            }
        }

        if !current_batch.is_empty() {
            batches.push(current_batch);
        }

        if batches.is_empty() {
            vec![Vec::new()]
        } else {
            batches
        }
    }

    /// Aggregate results from multiple LLM calls
    fn aggregate_results(&self, results: Vec<DispatchResult>) -> AggregatedResult {
        let text = match self.config.aggregation {
            AggregationStrategy::Concatenate => results
                .iter()
                .map(|r| r.text.as_str())
                .collect::<Vec<_>>()
                .join("\n\n"),
            AggregationStrategy::FirstOnly => {
                results.first().map(|r| r.text.clone()).unwrap_or_default()
            }
            AggregationStrategy::MajorityVote => {
                // For Phase 6, just concatenate (real voting logic deferred)
                results
                    .iter()
                    .map(|r| r.text.as_str())
                    .collect::<Vec<_>>()
                    .join("\n\n")
            }
        };

        AggregatedResult {
            text,
            total_input_tokens: results.iter().map(|r| r.input_tokens).sum(),
            total_output_tokens: results.iter().map(|r| r.output_tokens).sum(),
            num_calls: results.len(),
            total_duration_ms: results.iter().map(|r| r.duration_ms).sum(),
        }
    }

    /// Check if dispatcher has an LLM client configured
    pub fn is_configured(&self) -> bool {
        self.llm_client.is_some()
    }
}

#[cfg(test)]
mod tests {
    use async_trait::async_trait;
    use vapora_llm_router::providers::{CompletionResponse, ProviderError};

    use super::*;
    use crate::storage::Chunk;

    // Mock LLM client for testing
    struct MockLLMClient {
        response_text: String,
    }

    impl MockLLMClient {
        fn new(response: impl Into<String>) -> Self {
            Self {
                response_text: response.into(),
            }
        }
    }

    #[async_trait]
    impl LLMClient for MockLLMClient {
        async fn complete(
            &self,
            _prompt: String,
            _context: Option<String>,
        ) -> Result<CompletionResponse, ProviderError> {
            Ok(CompletionResponse {
                text: self.response_text.clone(),
                input_tokens: 100,
                output_tokens: 50,
                finish_reason: "stop".to_string(),
            })
        }

        async fn stream(
            &self,
            _prompt: String,
        ) -> Result<tokio::sync::mpsc::Receiver<String>, ProviderError> {
            let (tx, rx) = tokio::sync::mpsc::channel(1);
            let _ = tx.send(self.response_text.clone()).await;
            Ok(rx)
        }

        fn cost_per_1k_tokens(&self) -> f64 {
            0.001
        }

        fn latency_ms(&self) -> u32 {
            100
        }

        fn available(&self) -> bool {
            true
        }

        fn provider_name(&self) -> String {
            "mock".to_string()
        }

        fn model_name(&self) -> String {
            "mock-model".to_string()
        }

        fn calculate_cost(&self, input_tokens: u64, output_tokens: u64) -> u32 {
            ((input_tokens + output_tokens) as f64 * 0.001) as u32
        }
    }

    fn create_test_chunk(id: &str, content: &str) -> ScoredChunk {
        ScoredChunk {
            chunk: Chunk {
                chunk_id: id.to_string(),
                doc_id: "test-doc".to_string(),
                content: content.to_string(),
                embedding: None,
                start_idx: 0,
                end_idx: content.len(),
                metadata: None,
                created_at: chrono::Utc::now().to_rfc3339(),
            },
            score: 1.0,
            bm25_score: Some(0.8),
            semantic_score: Some(0.9),
        }
    }

    #[tokio::test]
    async fn test_dispatcher_creation() {
        let client = Arc::new(MockLLMClient::new("test response"));
        let dispatcher = LLMDispatcher::new(Some(client));
        assert!(dispatcher.is_configured());
    }

    #[tokio::test]
    async fn test_dispatcher_no_client() {
        let dispatcher = LLMDispatcher::new(None);
        assert!(!dispatcher.is_configured());

        let chunks = vec![create_test_chunk("chunk-1", "test content")];
        let result = dispatcher.dispatch("test query", &chunks).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_dispatch_single_chunk() {
        let client = Arc::new(MockLLMClient::new("Answer based on chunk"));
        let dispatcher = LLMDispatcher::new(Some(client));

        let chunks = vec![create_test_chunk("chunk-1", "Rust is awesome")];

        let result = dispatcher.dispatch("What is Rust?", &chunks).await.unwrap();
        assert_eq!(result.text, "Answer based on chunk");
        assert_eq!(result.num_calls, 1);
        assert_eq!(result.total_input_tokens, 100);
        assert_eq!(result.total_output_tokens, 50);
    }

    #[tokio::test]
    async fn test_dispatch_multiple_chunks() {
        let client = Arc::new(MockLLMClient::new("Combined answer"));
        let dispatcher = LLMDispatcher::new(Some(client));

        let chunks = vec![
            create_test_chunk("chunk-1", "Rust memory safety"),
            create_test_chunk("chunk-2", "Rust concurrency"),
            create_test_chunk("chunk-3", "Rust performance"),
        ];

        let result = dispatcher.dispatch("Explain Rust", &chunks).await.unwrap();
        assert_eq!(result.num_calls, 1); // All chunks in one batch (default max = 10)
        assert!(result.total_input_tokens > 0);
    }

    #[tokio::test]
    async fn test_batch_splitting() {
        let client = Arc::new(MockLLMClient::new("Batch response"));
        let config = DispatchConfig {
            max_chunks_per_dispatch: 2, // Small batch size
            ..Default::default()
        };

        let dispatcher = LLMDispatcher::with_config(Some(client), config);

        let chunks = vec![
            create_test_chunk("chunk-1", "content 1"),
            create_test_chunk("chunk-2", "content 2"),
            create_test_chunk("chunk-3", "content 3"),
            create_test_chunk("chunk-4", "content 4"),
            create_test_chunk("chunk-5", "content 5"),
        ];

        let result = dispatcher.dispatch("query", &chunks).await.unwrap();
        assert_eq!(result.num_calls, 3); // 5 chunks / 2 per batch = 3 batches
        assert_eq!(result.total_input_tokens, 300); // 3 calls * 100 tokens
        assert_eq!(result.total_output_tokens, 150); // 3 calls * 50 tokens
    }

    #[tokio::test]
    async fn test_aggregation_concatenate() {
        let client = Arc::new(MockLLMClient::new("Response part"));
        let config = DispatchConfig {
            max_chunks_per_dispatch: 1, // Force multiple calls
            aggregation: AggregationStrategy::Concatenate,
            ..Default::default()
        };

        let dispatcher = LLMDispatcher::with_config(Some(client), config);

        let chunks = vec![
            create_test_chunk("chunk-1", "content 1"),
            create_test_chunk("chunk-2", "content 2"),
        ];

        let result = dispatcher.dispatch("query", &chunks).await.unwrap();
        assert_eq!(result.text, "Response part\n\nResponse part");
    }

    #[tokio::test]
    async fn test_aggregation_first_only() {
        let client = Arc::new(MockLLMClient::new("First response"));
        let config = DispatchConfig {
            max_chunks_per_dispatch: 1,
            aggregation: AggregationStrategy::FirstOnly,
            ..Default::default()
        };

        let dispatcher = LLMDispatcher::with_config(Some(client), config);

        let chunks = vec![
            create_test_chunk("chunk-1", "content 1"),
            create_test_chunk("chunk-2", "content 2"),
        ];

        let result = dispatcher.dispatch("query", &chunks).await.unwrap();
        assert_eq!(result.text, "First response");
        assert_eq!(result.num_calls, 2); // Still makes all calls, just uses
                                         // first result
    }

    #[tokio::test]
    async fn test_empty_chunks() {
        let client = Arc::new(MockLLMClient::new("Response"));
        let dispatcher = LLMDispatcher::new(Some(client));

        let result = dispatcher.dispatch("query", &[]).await.unwrap();
        assert_eq!(result.num_calls, 0);
        assert_eq!(result.text, "");
    }

    #[tokio::test]
    async fn test_prompt_building() {
        let client = Arc::new(MockLLMClient::new("Response"));
        let dispatcher = LLMDispatcher::new(Some(client));

        let chunks = [create_test_chunk("chunk-1", "Test content")];
        let chunk_refs: Vec<&ScoredChunk> = chunks.iter().collect();

        let prompt = dispatcher.build_prompt("What is this?", &chunk_refs);

        assert!(prompt.contains("Query: What is this?"));
        assert!(prompt.contains("=== Chunk 1 ==="));
        assert!(prompt.contains("Test content"));
    }

    #[tokio::test]
    async fn test_prompt_with_metadata() {
        let client = Arc::new(MockLLMClient::new("Response"));
        let config = DispatchConfig {
            include_metadata: true,
            ..Default::default()
        };

        let dispatcher = LLMDispatcher::with_config(Some(client), config);

        let chunks = [create_test_chunk("chunk-1", "Content")];
        let chunk_refs: Vec<&ScoredChunk> = chunks.iter().collect();

        let prompt = dispatcher.build_prompt("query", &chunk_refs);

        assert!(prompt.contains("Metadata:"));
        assert!(prompt.contains("chunk_id=chunk-1"));
    }
}
