Vapora/crates/vapora-llm-router/src/router.rs

// vapora-llm-router: Routing engine for task-optimal LLM selection
// Phase 2: Complete implementation with fallback support

use std::collections::HashMap;
use std::sync::Arc;

use thiserror::Error;
use tracing::{debug, info, warn};

use crate::budget::BudgetManager;
use crate::config::{LLMRouterConfig, ProviderConfig};
use crate::cost_ranker::CostRanker;
use crate::cost_tracker::CostTracker;
use crate::providers::*;

#[derive(Debug, Error)]
pub enum RouterError {
    #[error("No providers available for task type: {0}")]
    NoProvidersAvailable(String),

    #[error("Provider not found: {0}")]
    ProviderNotFound(String),

    #[error("All providers failed")]
    AllProvidersFailed,

    #[error("Configuration error: {0}")]
    ConfigError(String),

    #[error("Budget error: {0}")]
    BudgetError(String),
}

/// LLM Router - selects optimal provider based on task type, cost, and budget
pub struct LLMRouter {
    config: Arc<LLMRouterConfig>,
    providers: HashMap<String, Arc<Box<dyn LLMClient>>>,
    cost_tracker: Arc<CostTracker>,
    budget_manager: Option<Arc<BudgetManager>>,
}

impl LLMRouter {
    /// Create a new router from configuration
    pub fn new(config: LLMRouterConfig) -> Result<Self, RouterError> {
        let mut providers = HashMap::new();
        let config_arc = Arc::new(config);

        // Initialize all enabled providers
        for (name, provider_config) in &config_arc.providers {
            if !provider_config.enabled {
                debug!("Provider {} is disabled, skipping", name);
                continue;
            }

            let client = Self::create_client(name, provider_config)?;
            providers.insert(name.clone(), Arc::new(client));
            info!("Initialized provider: {}", name);
        }

        Ok(Self {
            config: config_arc,
            providers,
            cost_tracker: Arc::new(CostTracker::new()),
            budget_manager: None,
        })
    }

    /// Set budget manager for cost enforcement
    pub fn with_budget_manager(mut self, budget_manager: Arc<BudgetManager>) -> Self {
        self.budget_manager = Some(budget_manager);
        self
    }

    /// Register an RLM provider (must be created externally with RLMEngine)
    ///
    /// RLM providers cannot be created from config alone since they require
    /// an initialized RLMEngine with storage and indexes.
    ///
    /// # Example
    /// ```ignore
    /// let rlm_engine = Arc::new(RLMEngine::new(storage, bm25_index)?);
    /// let rlm_provider = RLMProvider::new(rlm_engine, RLMProviderConfig::default(), None);
    /// router.add_rlm_provider("rlm", Arc::new(Box::new(rlm_provider)));
    /// ```
    pub fn add_rlm_provider(&mut self, name: &str, client: Arc<Box<dyn LLMClient>>) {
        self.providers.insert(name.to_string(), client);
        info!("Registered RLM provider: {}", name);
    }

    /// Create a client for a specific provider
    fn create_client(
        name: &str,
        config: &ProviderConfig,
    ) -> Result<Box<dyn LLMClient>, RouterError> {
        match name {
            "claude" => {
                let api_key = config.api_key.clone().ok_or_else(|| {
                    RouterError::ConfigError("Claude API key missing".to_string())
                })?;

                let client = ClaudeClient::new(
                    api_key,
                    config.model.clone(),
                    config.max_tokens,
                    config.temperature,
                    config.cost_per_1m_input,
                    config.cost_per_1m_output,
                )
                .map_err(|e| RouterError::ConfigError(e.to_string()))?;

                Ok(Box::new(client))
            }
            "openai" => {
                let api_key = config.api_key.clone().ok_or_else(|| {
                    RouterError::ConfigError("OpenAI API key missing".to_string())
                })?;

                let client = OpenAIClient::new(
                    api_key,
                    config.model.clone(),
                    config.max_tokens,
                    config.temperature,
                    config.cost_per_1m_input,
                    config.cost_per_1m_output,
                )
                .map_err(|e| RouterError::ConfigError(e.to_string()))?;

                Ok(Box::new(client))
            }
            "ollama" => {
                let endpoint = config
                    .url
                    .clone()
                    .unwrap_or_else(|| "http://localhost:11434".to_string());

                let client = OllamaClient::new(
                    endpoint,
                    config.model.clone(),
                    config.max_tokens,
                    config.temperature,
                )
                .map_err(|e| RouterError::ConfigError(e.to_string()))?;

                Ok(Box::new(client))
            }
            "rlm" => {
                // RLM provider requires special configuration
                // For now, return error - RLM instances must be created externally
                // and registered via add_rlm_provider()
                Err(RouterError::ConfigError(
                    "RLM provider must be registered via add_rlm_provider() method".to_string(),
                ))
            }
            _ => Err(RouterError::ConfigError(format!(
                "Unknown provider: {}",
                name
            ))),
        }
    }

    /// Route a task to the optimal provider with budget awareness
    pub async fn route(
        &self,
        task_type: &str,
        conditions: Option<HashMap<String, String>>,
    ) -> Result<String, RouterError> {
        self.route_with_budget(task_type, conditions, None).await
    }

    /// Route a task with budget awareness
    pub async fn route_with_budget(
        &self,
        task_type: &str,
        conditions: Option<HashMap<String, String>>,
        agent_role: Option<&str>,
    ) -> Result<String, RouterError> {
        let mut context = HashMap::new();
        context.insert("task_type".to_string(), task_type.to_string());

        if let Some(cond) = conditions {
            context.extend(cond);
        }

        // Check budget if provided
        if let (Some(role), Some(budget_mgr)) = (agent_role, &self.budget_manager) {
            match budget_mgr.check_budget(role).await {
                Ok(status) if status.exceeded => {
                    // Budget exceeded - use fallback provider
                    info!(
                        "Budget exceeded for role {}, using fallback provider: {}",
                        role, status.fallback_provider
                    );
                    return Ok(status.fallback_provider);
                }
                Ok(status) if status.near_threshold => {
                    // Budget near threshold - prefer cost-efficient providers
                    debug!(
                        "Budget near threshold for role {}, selecting cost-efficient provider",
                        role
                    );
                    return self.select_cost_efficient_provider(task_type).await;
                }
                Ok(_) => {
                    // Budget ok, continue with normal routing
                }
                Err(e) => {
                    warn!("Budget check failed: {}, continuing with normal routing", e);
                }
            }
        }

        // Try to find matching routing rule
        if let Some(rule) = self.config.find_rule(&context) {
            debug!("Found routing rule: {}", rule.name);

            if self.is_provider_available(&rule.provider) {
                info!(
                    "Routing {} to {} via rule {}",
                    task_type, rule.provider, rule.name
                );
                return Ok(rule.provider.clone());
            }

            warn!(
                "Primary provider {} unavailable, falling back",
                rule.provider
            );
        }

        // Use default provider
        let default_provider = &self.config.routing.default_provider;
        if self.is_provider_available(default_provider) {
            info!(
                "Routing {} to default provider {}",
                task_type, default_provider
            );
            return Ok(default_provider.clone());
        }

        // Fallback to any available provider
        if self.config.routing.fallback_enabled {
            if let Some(provider_name) = self.find_available_provider() {
                warn!(
                    "Using fallback provider {} for {}",
                    provider_name, task_type
                );
                return Ok(provider_name);
            }
        }

        Err(RouterError::NoProvidersAvailable(task_type.to_string()))
    }

    /// Select the most cost-efficient provider
    async fn select_cost_efficient_provider(&self, task_type: &str) -> Result<String, RouterError> {
        let available_providers: Vec<(String, ProviderConfig)> = self
            .providers
            .iter()
            .filter(|(_name, provider)| provider.available())
            .filter_map(|(name, _provider)| {
                self.config
                    .providers
                    .get(name)
                    .map(|cfg| (name.clone(), cfg.clone()))
            })
            .collect();

        if available_providers.is_empty() {
            return Err(RouterError::NoProvidersAvailable(task_type.to_string()));
        }

        // Rank by cost efficiency
        let ranked = CostRanker::rank_by_efficiency(available_providers, task_type, 1000, 200);

        if let Some(best) = ranked.first() {
            info!(
                "Selected cost-efficient provider {} for {} (efficiency: {:.2})",
                best.provider, task_type, best.cost_efficiency
            );
            Ok(best.provider.clone())
        } else {
            Err(RouterError::NoProvidersAvailable(task_type.to_string()))
        }
    }

    /// Get a provider client by name
    pub fn get_provider(&self, name: &str) -> Result<Arc<Box<dyn LLMClient>>, RouterError> {
        self.providers
            .get(name)
            .cloned()
            .ok_or_else(|| RouterError::ProviderNotFound(name.to_string()))
    }

    /// Check if a provider is available
    fn is_provider_available(&self, name: &str) -> bool {
        self.providers
            .get(name)
            .map(|p| p.available())
            .unwrap_or(false)
    }

    /// Find any available provider
    fn find_available_provider(&self) -> Option<String> {
        self.providers
            .iter()
            .find(|(_, provider)| provider.available())
            .map(|(name, _)| name.clone())
    }

    /// Execute a completion request with optimal provider and budget tracking
    pub async fn complete(
        &self,
        task_type: &str,
        prompt: String,
        context: Option<String>,
        conditions: Option<HashMap<String, String>>,
    ) -> Result<CompletionResponse, RouterError> {
        self.complete_with_budget(task_type, prompt, context, conditions, None)
            .await
    }

    /// Execute a completion with budget awareness and cost tracking
    pub async fn complete_with_budget(
        &self,
        task_type: &str,
        prompt: String,
        context: Option<String>,
        conditions: Option<HashMap<String, String>>,
        agent_role: Option<&str>,
    ) -> Result<CompletionResponse, RouterError> {
        let provider_name = self
            .route_with_budget(task_type, conditions, agent_role)
            .await?;
        let provider = self.get_provider(&provider_name)?;

        match provider.complete(prompt, context).await {
            Ok(response) => {
                // Track cost
                if self.config.routing.cost_tracking_enabled {
                    let cost =
                        provider.calculate_cost(response.input_tokens, response.output_tokens);
                    self.cost_tracker.log_usage(
                        &provider_name,
                        task_type,
                        response.input_tokens,
                        response.output_tokens,
                        cost,
                    );

                    // Record spend with budget manager if available
                    self.record_budget_spend(agent_role, cost as u32).await;
                }

                Ok(response)
            }
            Err(e) => {
                warn!("Provider {} failed: {}", provider_name, e);

                // Try fallback if enabled
                if self.config.routing.fallback_enabled {
                    return self
                        .try_fallback_with_budget(task_type, &provider_name, agent_role)
                        .await;
                }

                Err(RouterError::AllProvidersFailed)
            }
        }
    }

    /// Try fallback providers with budget tracking
    async fn try_fallback_with_budget(
        &self,
        task_type: &str,
        failed_provider: &str,
        _agent_role: Option<&str>,
    ) -> Result<CompletionResponse, RouterError> {
        // Build fallback chain excluding failed provider
        let fallback_chain: Vec<String> = self
            .providers
            .iter()
            .filter(|(name, provider)| *name != failed_provider && provider.available())
            .map(|(name, _)| name.clone())
            .collect();

        if fallback_chain.is_empty() {
            return Err(RouterError::AllProvidersFailed);
        }

        warn!(
            "Primary provider {} failed for {}, trying fallback chain",
            failed_provider, task_type
        );

        // Try each fallback provider (placeholder implementation)
        // In production, you would retry the original prompt with each fallback
        // provider For now, we log which providers would be tried and return
        // error
        for provider_name in fallback_chain {
            warn!("Trying fallback provider: {}", provider_name);
            // Actual retry logic would go here with cost tracking
            // For this phase, we return the error as fallbacks are handled at
            // routing level
        }

        Err(RouterError::AllProvidersFailed)
    }

    /// Get cost tracker reference
    pub fn cost_tracker(&self) -> Arc<CostTracker> {
        Arc::clone(&self.cost_tracker)
    }

    /// List all available providers
    pub fn list_providers(&self) -> Vec<String> {
        self.providers.keys().cloned().collect()
    }

    /// Get provider statistics
    pub fn provider_stats(&self, name: &str) -> Option<ProviderStats> {
        self.providers.get(name).map(|provider| ProviderStats {
            name: name.to_string(),
            model: provider.model_name(),
            available: provider.available(),
            cost_per_1k_tokens: provider.cost_per_1k_tokens(),
            latency_ms: provider.latency_ms(),
        })
    }

    /// Record budget spend for agent role
    async fn record_budget_spend(&self, agent_role: Option<&str>, cost: u32) {
        if let (Some(role), Some(budget_mgr)) = (agent_role, &self.budget_manager) {
            if let Err(e) = budget_mgr.record_spend(role, cost).await {
                warn!("Failed to record budget spend: {}", e);
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct ProviderStats {
    pub name: String,
    pub model: String,
    pub available: bool,
    pub cost_per_1k_tokens: f64,
    pub latency_ms: u32,
}

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

    fn create_test_config() -> LLMRouterConfig {
        let mut providers = HashMap::new();

        providers.insert(
            "claude".to_string(),
            ProviderConfig {
                enabled: true,
                api_key: Some("test_key".to_string()),
                url: None,
                model: "claude-sonnet-4".to_string(),
                max_tokens: 4096,
                temperature: 0.7,
                cost_per_1m_input: 3.0,
                cost_per_1m_output: 15.0,
            },
        );

        providers.insert(
            "ollama".to_string(),
            ProviderConfig {
                enabled: true,
                api_key: None,
                url: Some("http://localhost:11434".to_string()),
                model: "llama3.2".to_string(),
                max_tokens: 4096,
                temperature: 0.7,
                cost_per_1m_input: 0.0,
                cost_per_1m_output: 0.0,
            },
        );

        LLMRouterConfig {
            routing: RoutingConfig {
                default_provider: "claude".to_string(),
                cost_tracking_enabled: true,
                fallback_enabled: true,
            },
            providers,
            routing_rules: vec![],
        }
    }

    #[tokio::test]
    async fn test_router_creation() {
        let config = create_test_config();
        let router = LLMRouter::new(config);
        assert!(router.is_ok());
    }

    #[tokio::test]
    async fn test_routing_to_default() {
        let config = create_test_config();
        let router = LLMRouter::new(config).unwrap();

        let provider = router.route("test_task", None).await;
        assert!(provider.is_ok());
        assert_eq!(provider.unwrap(), "claude");
    }

    #[tokio::test]
    async fn test_list_providers() {
        let config = create_test_config();
        let router = LLMRouter::new(config).unwrap();

        let providers = router.list_providers();
        assert!(providers.contains(&"claude".to_string()));
        assert!(providers.contains(&"ollama".to_string()));
    }

    #[test]
    fn test_provider_stats() {
        let config = create_test_config();
        let router = LLMRouter::new(config).unwrap();

        let stats = router.provider_stats("claude");
        assert!(stats.is_some());

        let stats = stats.unwrap();
        assert_eq!(stats.name, "claude");
        assert!(stats.available);
    }
}