//! Multi-IA router for selecting optimal LLM provider
//!
//! Routes tasks to the best LLM provider based on cost, capability, and availability.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Supported LLM providers
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Serialize, Deserialize)]
pub enum LLMProvider {
    /// Anthropic Claude
    Claude,
    /// OpenAI GPT models
    OpenAI,
    /// Google Gemini
    Gemini,
    /// Open source Ollama
    Ollama,
}

impl LLMProvider {
    /// Get provider name
    pub fn name(&self) -> &'static str {
        match self {
            Self::Claude => "claude",
            Self::OpenAI => "openai",
            Self::Gemini => "gemini",
            Self::Ollama => "ollama",
        }
    }

    /// Get default model for provider
    pub fn default_model(&self) -> &'static str {
        match self {
            Self::Claude => "claude-sonnet-4.5",
            Self::OpenAI => "gpt-4",
            Self::Gemini => "gemini-pro",
            Self::Ollama => "mistral",
        }
    }

    /// Estimated cost per 1M tokens (in USD)
    pub fn cost_per_1m_tokens(&self) -> f64 {
        match self {
            Self::Claude => 3.00,  // Claude Sonnet approximate
            Self::OpenAI => 15.00, // GPT-4 approximate
            Self::Gemini => 2.50,  // Gemini approximate
            Self::Ollama => 0.00,  // Self-hosted
        }
    }

    /// Capability score (0-100)
    pub fn capability_score(&self) -> u32 {
        match self {
            Self::Claude => 95,
            Self::OpenAI => 90,
            Self::Gemini => 85,
            Self::Ollama => 70,
        }
    }
}

/// LLM routing strategy
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
pub enum RoutingStrategy {
    /// Use cheapest provider
    Cost,
    /// Use most capable provider
    Capability,
    /// Balance cost and capability
    Balanced,
    /// Round-robin across providers
    RoundRobin,
    /// Use provider specified in context
    Manual,
}

/// Routing decision
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingDecision {
    /// Selected provider
    pub provider: LLMProvider,
    /// Selected model
    pub model: String,
    /// Routing reason
    pub reason: String,
    /// Estimated cost
    pub estimated_cost: f64,
    /// Fallback providers if primary fails
    pub fallbacks: Vec<LLMProvider>,
}

/// Multi-IA router
#[derive(Debug, Clone)]
pub struct MultiIARouter {
    /// Provider status map
    providers: HashMap<LLMProvider, ProviderStatus>,
}

/// Provider availability status
#[derive(Debug, Clone)]
struct ProviderStatus {
    available: bool,
    latency_ms: u32,
    error_rate: f64,
    requests_today: u32,
}

impl MultiIARouter {
    /// Create new router with default providers
    pub fn new() -> Self {
        let mut providers = HashMap::new();

        for provider in [
            LLMProvider::Claude,
            LLMProvider::OpenAI,
            LLMProvider::Gemini,
            LLMProvider::Ollama,
        ] {
            providers.insert(
                provider,
                ProviderStatus {
                    available: true,
                    latency_ms: 500,
                    error_rate: 0.01,
                    requests_today: 0,
                },
            );
        }

        Self { providers }
    }

    /// Route request to best provider
    pub fn route(
        &self,
        strategy: RoutingStrategy,
        context: &HashMap<String, String>,
    ) -> RoutingDecision {
        match strategy {
            RoutingStrategy::Cost => self.route_by_cost(),
            RoutingStrategy::Capability => self.route_by_capability(),
            RoutingStrategy::Balanced => self.route_balanced(),
            RoutingStrategy::RoundRobin => self.route_round_robin(),
            RoutingStrategy::Manual => self.route_manual(context),
        }
    }

    fn route_by_cost(&self) -> RoutingDecision {
        let provider = LLMProvider::Ollama; // Cheapest
        RoutingDecision {
            provider,
            model: provider.default_model().to_string(),
            reason: "Selected for lowest cost".to_string(),
            estimated_cost: provider.cost_per_1m_tokens(),
            fallbacks: vec![
                LLMProvider::Gemini,
                LLMProvider::Claude,
                LLMProvider::OpenAI,
            ],
        }
    }

    fn route_by_capability(&self) -> RoutingDecision {
        let provider = LLMProvider::Claude; // Most capable
        RoutingDecision {
            provider,
            model: provider.default_model().to_string(),
            reason: "Selected for highest capability".to_string(),
            estimated_cost: provider.cost_per_1m_tokens(),
            fallbacks: vec![
                LLMProvider::OpenAI,
                LLMProvider::Gemini,
                LLMProvider::Ollama,
            ],
        }
    }

    fn route_balanced(&self) -> RoutingDecision {
        let provider = LLMProvider::Claude; // Good balance
        RoutingDecision {
            provider,
            model: provider.default_model().to_string(),
            reason: "Selected for balanced cost and capability".to_string(),
            estimated_cost: provider.cost_per_1m_tokens(),
            fallbacks: vec![
                LLMProvider::Gemini,
                LLMProvider::OpenAI,
                LLMProvider::Ollama,
            ],
        }
    }

    fn route_round_robin(&self) -> RoutingDecision {
        // Simple round-robin implementation
        let provider = LLMProvider::Claude;
        RoutingDecision {
            provider,
            model: provider.default_model().to_string(),
            reason: "Selected for round-robin distribution".to_string(),
            estimated_cost: provider.cost_per_1m_tokens(),
            fallbacks: vec![
                LLMProvider::OpenAI,
                LLMProvider::Gemini,
                LLMProvider::Ollama,
            ],
        }
    }

    fn route_manual(&self, context: &HashMap<String, String>) -> RoutingDecision {
        let provider_name = context
            .get("provider")
            .map(|s| s.as_str())
            .unwrap_or("claude");

        let provider = match provider_name {
            "openai" => LLMProvider::OpenAI,
            "gemini" => LLMProvider::Gemini,
            "ollama" => LLMProvider::Ollama,
            _ => LLMProvider::Claude,
        };

        RoutingDecision {
            provider,
            model: provider.default_model().to_string(),
            reason: "Selected by manual specification".to_string(),
            estimated_cost: provider.cost_per_1m_tokens(),
            fallbacks: vec![LLMProvider::Claude, LLMProvider::OpenAI],
        }
    }

    /// Get provider health status
    pub fn get_provider_status(&self, provider: LLMProvider) -> Option<String> {
        self.providers.get(&provider).map(|status| {
            if status.available {
                format!(
                    "{}: Available (latency: {}ms, error_rate: {:.2}%)",
                    provider.name(),
                    status.latency_ms,
                    status.error_rate * 100.0
                )
            } else {
                format!("{}: Unavailable", provider.name())
            }
        })
    }

    /// Mark provider as unavailable
    pub fn mark_unavailable(&mut self, provider: LLMProvider) {
        if let Some(status) = self.providers.get_mut(&provider) {
            status.available = false;
        }
    }

    /// Get router statistics
    pub fn get_statistics(&self) -> RouterStats {
        let total_requests: u32 = self.providers.values().map(|s| s.requests_today).sum();
        let providers_available = self.providers.values().filter(|s| s.available).count();

        RouterStats {
            total_requests,
            providers_available,
            providers_total: self.providers.len(),
        }
    }
}

impl Default for MultiIARouter {
    fn default() -> Self {
        Self::new()
    }
}

/// Router statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterStats {
    /// Total requests today
    pub total_requests: u32,
    /// Available providers
    pub providers_available: usize,
    /// Total providers
    pub providers_total: usize,
}

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

    #[test]
    fn test_llm_provider_names() {
        assert_eq!(LLMProvider::Claude.name(), "claude");
        assert_eq!(LLMProvider::OpenAI.name(), "openai");
    }

    #[test]
    fn test_provider_models() {
        assert!(!LLMProvider::Claude.default_model().is_empty());
        assert!(!LLMProvider::OpenAI.default_model().is_empty());
    }

    #[test]
    fn test_provider_costs() {
        assert!(LLMProvider::Claude.cost_per_1m_tokens() > 0.0);
        assert_eq!(LLMProvider::Ollama.cost_per_1m_tokens(), 0.0);
    }

    #[test]
    fn test_capability_scores() {
        assert!(LLMProvider::Claude.capability_score() > LLMProvider::Ollama.capability_score());
    }

    #[test]
    fn test_router_creation() {
        let router = MultiIARouter::new();
        assert_eq!(router.providers.len(), 4);
    }

    #[test]
    fn test_route_by_cost() {
        let router = MultiIARouter::new();
        let decision = router.route_by_cost();
        assert_eq!(decision.provider, LLMProvider::Ollama);
    }

    #[test]
    fn test_route_by_capability() {
        let router = MultiIARouter::new();
        let decision = router.route_by_capability();
        assert_eq!(decision.provider, LLMProvider::Claude);
    }

    #[test]
    fn test_route_balanced() {
        let router = MultiIARouter::new();
        let decision = router.route_balanced();
        assert!(
            decision.provider == LLMProvider::Claude || decision.provider == LLMProvider::Gemini
        );
    }

    #[test]
    fn test_manual_routing() {
        let router = MultiIARouter::new();
        let mut context = HashMap::new();
        context.insert("provider".to_string(), "openai".to_string());

        let decision = router.route_manual(&context);
        assert_eq!(decision.provider, LLMProvider::OpenAI);
    }

    #[test]
    fn test_provider_unavailable() {
        let mut router = MultiIARouter::new();
        router.mark_unavailable(LLMProvider::Claude);

        let status = router.get_provider_status(LLMProvider::Claude);
        assert!(status.unwrap().contains("Unavailable"));
    }

    #[test]
    fn test_router_statistics() {
        let router = MultiIARouter::new();
        let stats = router.get_statistics();
        assert_eq!(stats.providers_total, 4);
    }
}
