// RLM Provider - LLMClient implementation for LLMRouter integration
// Routes long-context tasks to RLM engine, short tasks to fallback LLM

use std::sync::Arc;

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

use crate::storage::Storage;
use crate::RLMEngine;

/// RLM Provider configuration
#[derive(Debug, Clone)]
pub struct RLMProviderConfig {
    /// Context length threshold - use RLM for prompts longer than this (in
    /// characters)
    pub context_threshold: usize,
    /// Number of chunks to retrieve for hybrid search
    pub top_k_chunks: usize,
    /// Whether to enable distributed LLM dispatch (future feature)
    pub enable_llm_dispatch: bool,
}

impl Default for RLMProviderConfig {
    fn default() -> Self {
        Self {
            context_threshold: 50_000, // 50k characters (~12.5k tokens)
            top_k_chunks: 5,
            enable_llm_dispatch: false, // Not implemented yet
        }
    }
}

/// RLM Provider - implements LLMClient for LLMRouter integration
///
/// Routes tasks based on context length:
/// - Short contexts (< threshold): Fallback to standard LLM provider
/// - Long contexts (> threshold): Use RLM chunking + hybrid search +
///   distributed reasoning
///
/// # Example
/// ```ignore
/// use vapora_rlm::{RLMEngine, RLMProvider, RLMProviderConfig};
///
/// let engine = Arc::new(RLMEngine::new(storage, bm25_index)?);
/// let config = RLMProviderConfig::default();
/// let rlm_provider = RLMProvider::new(engine, config, Some(fallback_llm));
///
/// // Register with LLMRouter
/// router.add_rlm_provider("rlm", Arc::new(Box::new(rlm_provider)));
/// ```
pub struct RLMProvider<S: Storage> {
    engine: Arc<RLMEngine<S>>,
    config: RLMProviderConfig,
    fallback_client: Option<Arc<Box<dyn LLMClient>>>,
}

impl<S: Storage + 'static> RLMProvider<S> {
    /// Create a new RLM provider
    ///
    /// # Arguments
    /// - `engine`: RLMEngine instance (manages chunking, storage, hybrid
    ///   search)
    /// - `config`: RLMProvider configuration
    /// - `fallback_client`: Optional LLM client for short contexts (<
    ///   threshold)
    pub fn new(
        engine: Arc<RLMEngine<S>>,
        config: RLMProviderConfig,
        fallback_client: Option<Arc<Box<dyn LLMClient>>>,
    ) -> Self {
        Self {
            engine,
            config,
            fallback_client,
        }
    }

    /// Check if context should use RLM (based on length threshold)
    fn should_use_rlm(&self, prompt: &str, context: Option<&str>) -> bool {
        let total_length = prompt.len() + context.map(|c| c.len()).unwrap_or(0);
        let use_rlm = total_length > self.config.context_threshold;

        debug!(
            "Context length: {}, threshold: {}, using RLM: {}",
            total_length, self.config.context_threshold, use_rlm
        );

        use_rlm
    }

    /// Generate a unique document ID for this request
    fn generate_doc_id(&self) -> String {
        format!("rlm-{}", uuid::Uuid::new_v4())
    }
}

#[async_trait]
impl<S: Storage + 'static> LLMClient for RLMProvider<S> {
    async fn complete(
        &self,
        prompt: String,
        context: Option<String>,
    ) -> Result<CompletionResponse, ProviderError> {
        // Decide: RLM or fallback?
        if self.should_use_rlm(&prompt, context.as_deref()) {
            info!("Using RLM for long-context task");

            let doc_id = self.generate_doc_id();

            // Combine prompt + context into document
            let content = if let Some(ctx) = context {
                format!("{}\n\n{}", ctx, prompt)
            } else {
                prompt.clone()
            };

            // Load document (chunk + embed + index)
            let _chunk_count = self
                .engine
                .load_document(&doc_id, &content, None)
                .await
                .map_err(|e| ProviderError::LlmError(format!("RLM load failed: {}", e)))?;

            // Query with hybrid search (BM25 + semantic + RRF)
            let results = self
                .engine
                .query(&doc_id, &prompt, None, self.config.top_k_chunks)
                .await
                .map_err(|e| ProviderError::LlmError(format!("RLM query failed: {}", e)))?;

            // Aggregate chunks into response
            // Note: Future enhancement - dispatch chunks to LLM for synthesis
            let text = results
                .iter()
                .map(|r| r.chunk.content.as_str())
                .collect::<Vec<_>>()
                .join("\n\n");

            let text_len = text.len();
            debug!(
                "RLM retrieved {} chunks, total length: {}",
                results.len(),
                text_len
            );

            Ok(CompletionResponse {
                text: text.clone(),
                input_tokens: content.len() as u64 / 4, // Rough estimate: 4 chars/token
                output_tokens: text_len as u64 / 4,
                finish_reason: "rlm_retrieval".to_string(),
            })
        } else {
            // Short context - fallback to standard LLM
            debug!("Using fallback LLM for short-context task");

            if let Some(fallback) = &self.fallback_client {
                fallback.complete(prompt, context).await
            } else {
                Err(ProviderError::ConfigError(
                    "RLM fallback provider not configured for short contexts".to_string(),
                ))
            }
        }
    }

    async fn stream(
        &self,
        prompt: String,
    ) -> Result<tokio::sync::mpsc::Receiver<String>, ProviderError> {
        // Streaming not implemented for RLM yet
        // Fallback to complete() and stream the result
        let response = self.complete(prompt, None).await?;

        let (tx, rx) = tokio::sync::mpsc::channel(1);
        tokio::spawn(async move {
            let _ = tx.send(response.text).await;
        });

        Ok(rx)
    }

    fn cost_per_1k_tokens(&self) -> f64 {
        // RLM cost is primarily storage + index, minimal LLM usage
        // Retrieval is much cheaper than generation
        if let Some(fallback) = &self.fallback_client {
            fallback.cost_per_1k_tokens() * 0.1 // 10% of fallback cost
                                                // (retrieval vs generation)
        } else {
            0.01 // 1 cent per 1k tokens (storage cost estimate)
        }
    }

    fn latency_ms(&self) -> u32 {
        // RLM target: <500ms for load + query workflow
        // (from performance_test.rs target)
        500
    }

    fn available(&self) -> bool {
        // RLM is always available if engine is initialized
        true
    }

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

    fn model_name(&self) -> String {
        "rlm-hybrid-search".to_string()
    }

    fn calculate_cost(&self, input_tokens: u64, output_tokens: u64) -> u32 {
        // Simple cost model: primarily storage + indexing cost
        // Much cheaper than LLM generation
        let storage_cost = ((input_tokens + output_tokens) as f64 / 1_000_000.0) * 0.1; // 10 cents per 1M tokens
        (storage_cost * 100.0) as u32 // Convert to cents
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::search::bm25::BM25Index;
    use crate::storage::MockStorage;

    #[tokio::test]
    async fn test_rlm_provider_context_routing() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());
        let engine = Arc::new(RLMEngine::new(storage, bm25_index).unwrap());

        let config = RLMProviderConfig {
            context_threshold: 100, // Very low for testing
            top_k_chunks: 5,
            enable_llm_dispatch: false,
        };

        let provider = RLMProvider::new(engine, config, None);

        // Test short context detection
        assert!(!provider.should_use_rlm("short prompt", None));

        // Test long context detection
        let long_prompt = "x".repeat(150);
        assert!(provider.should_use_rlm(&long_prompt, None));

        // Test context + prompt combination
        assert!(provider.should_use_rlm("short", Some("x".repeat(100).as_str())));
    }

    #[tokio::test]
    async fn test_rlm_provider_long_context_complete() {
        let storage = Arc::new(MockStorage::new());
        let bm25_index = Arc::new(BM25Index::new().unwrap());
        let engine = Arc::new(RLMEngine::new(storage, bm25_index).unwrap());

        let config = RLMProviderConfig {
            context_threshold: 10, // Very low for testing
            top_k_chunks: 5,
            enable_llm_dispatch: false,
        };

        let provider = RLMProvider::new(engine, config, None);

        // Test with long context
        let long_content = "This is a test document with multiple lines.\n".repeat(100);
        let result = provider
            .complete("What is this about?".to_string(), Some(long_content))
            .await;

        assert!(result.is_ok());
        let response = result.unwrap();
        assert!(!response.text.is_empty());
        assert_eq!(response.finish_reason, "rlm_retrieval");
    }
}
