//! LLM Provider Router (Phase-78: T78.8)
//!
//! Routes chat requests to the optimal provider/model, handles failover,
//! and manages provider health state.
//!
//! # Architecture
//!
//! ```text
//! ┌──────────────────────────────────────────────┐
//! │              ProviderRouter                   │
//! │  ┌─────────────────────────────────────────┐ │
//! │  │ route(model_name, config) → LlmClient   │ │
//! │  └─────────┬───────────────────────────────┘ │
//! │  ┌─────────┴─────────────────────────────────┤
//! │  │  Health Monitor   │  Strategy Selector    │ │
//! │  │  (periodic check) │  (priority/round-     │ │
//! │  │                   │   robin/cheapest)     │ │
//! │  └───────────────────┴───────────────────────┘
//! └──────────────────────────────────────────────┘
//! ```
//!
//! # Routing strategies
//!
//! - **Priority**: Use the first healthy provider (in registration order)
//! - **RoundRobin**: Cycle through healthy providers
//! - **Cheapest**: Select the cheapest provider for the model

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicBool, Ordering};

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

use super::client::LlmClient;
use super::error::{LlmError, LlmResult};

// ============================================================================
// Routing strategy
// ============================================================================

/// How to select among multiple providers that serve the same model.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingStrategy {
    /// Use the first healthy provider (ordered by priority).
    Priority,
    /// Round-robin across healthy providers.
    RoundRobin,
    /// Select the cheapest provider for the model.
    Cheapest,
}

impl Default for RoutingStrategy {
    fn default() -> Self {
        Self::Priority
    }
}

// ============================================================================
// Provider health state
// ============================================================================

/// Tracks health state for a single provider.
struct ProviderHealth {
    healthy: AtomicBool,
    last_check_ms: AtomicU64,
    consecutive_failures: AtomicU64,
}

impl ProviderHealth {
    fn new() -> Self {
        Self {
            healthy: AtomicBool::new(true), // Assume healthy until proven otherwise
            last_check_ms: AtomicU64::new(0),
            consecutive_failures: AtomicU64::new(0),
        }
    }

    fn is_healthy(&self) -> bool {
        self.healthy.load(Ordering::Relaxed)
    }

    fn mark_healthy(&self) {
        self.healthy.store(true, Ordering::Relaxed);
        self.consecutive_failures.store(0, Ordering::Relaxed);
        self.last_check_ms.store(now_ms(), Ordering::Relaxed);
    }

    fn mark_unhealthy(&self) {
        self.healthy.store(false, Ordering::Relaxed);
        self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
        self.last_check_ms.store(now_ms(), Ordering::Relaxed);
    }

    /// Should we recheck this provider? (After cool-down period)
    fn should_recheck(&self, cooldown_ms: u64) -> bool {
        let last = self.last_check_ms.load(Ordering::Relaxed);
        let now = now_ms();
        now.saturating_sub(last) >= cooldown_ms
    }
}

// ============================================================================
// ProviderRouter
// ============================================================================

/// The main router that manages LLM clients and routes requests.
pub struct ProviderRouter {
    /// Registered clients by provider ID.
    clients: DashMap<String, Arc<dyn LlmClient>>,

    /// Model → [provider_id, ...] mapping (which providers serve which models).
    model_providers: DashMap<String, Vec<String>>,

    /// Provider health state by provider ID.
    health: DashMap<String, Arc<ProviderHealth>>,

    /// Default routing strategy.
    strategy: RoutingStrategy,

    /// Round-robin counter (global).
    rr_counter: AtomicU64,

    /// Health check cool-down (milliseconds).
    health_cooldown_ms: u64,
}

impl ProviderRouter {
    /// Create a new router with the given strategy.
    pub fn new(strategy: RoutingStrategy) -> Self {
        Self {
            clients: DashMap::new(),
            model_providers: DashMap::new(),
            health: DashMap::new(),
            strategy,
            rr_counter: AtomicU64::new(0),
            health_cooldown_ms: 30_000, // 30 seconds default
        }
    }

    /// Create with default (Priority) strategy.
    pub fn default_router() -> Self {
        Self::new(RoutingStrategy::Priority)
    }

    /// Register a provider client.
    pub fn register_client(
        &self,
        provider_id: impl Into<String>,
        client: Arc<dyn LlmClient>,
    ) {
        let id = provider_id.into();
        debug!(provider_id = %id, "Registering LLM client");
        self.health
            .entry(id.clone())
            .or_insert_with(|| Arc::new(ProviderHealth::new()));
        self.clients.insert(id, client);
    }

    /// Register that a provider serves a specific model.
    pub fn register_model_provider(
        &self,
        model_name: impl Into<String>,
        provider_id: impl Into<String>,
    ) {
        let model = model_name.into();
        let pid = provider_id.into();
        self.model_providers
            .entry(model)
            .or_insert_with(Vec::new)
            .push(pid);
    }

    /// Unregister a provider.
    pub fn unregister_provider(&self, provider_id: &str) {
        self.clients.remove(provider_id);
        self.health.remove(provider_id);
        // Remove from model_providers
        self.model_providers.iter_mut().for_each(|mut entry| {
            entry.value_mut().retain(|pid| pid != provider_id);
        });
    }

    /// Get a client by provider ID (direct access, no routing).
    pub fn get_client(&self, provider_id: &str) -> Option<Arc<dyn LlmClient>> {
        self.clients.get(provider_id).map(|c| c.value().clone())
    }

    /// Route to the best available client for a given model.
    ///
    /// Returns the selected client and the provider ID.
    pub fn route(&self, model_name: &str) -> LlmResult<(Arc<dyn LlmClient>, String)> {
        let provider_ids = self
            .model_providers
            .get(model_name)
            .map(|v| v.value().clone())
            .unwrap_or_default();

        if provider_ids.is_empty() {
            return Err(LlmError::ModelNotFound(format!(
                "No provider registered for model '{model_name}'"
            )));
        }

        // Filter to healthy providers
        let healthy_providers: Vec<&String> = provider_ids
            .iter()
            .filter(|pid| {
                self.health
                    .get(*pid)
                    .map(|h| h.is_healthy())
                    .unwrap_or(false)
            })
            .collect();

        let candidates = if healthy_providers.is_empty() {
            // All providers are unhealthy — check if any should be rechecked
            let rechecks: Vec<&String> = provider_ids
                .iter()
                .filter(|pid| {
                    self.health
                        .get(*pid)
                        .map(|h| h.should_recheck(self.health_cooldown_ms))
                        .unwrap_or(true)
                })
                .collect();

            if rechecks.is_empty() {
                // All providers are unhealthy and within cool-down
                return Err(LlmError::AllProvidersUnavailable(format!(
                    "All providers for model '{model_name}' are unavailable"
                )));
            }
            rechecks
        } else {
            healthy_providers
        };

        // Select based on strategy
        let selected_id = match self.strategy {
            RoutingStrategy::Priority => {
                // First in the list
                candidates[0]
            }
            RoutingStrategy::RoundRobin => {
                let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed) as usize;
                candidates[idx % candidates.len()]
            }
            RoutingStrategy::Cheapest => {
                // For now, same as Priority (TODO: integrate with model pricing)
                candidates[0]
            }
        };

        let client = self
            .clients
            .get(selected_id)
            .map(|c| c.value().clone())
            .ok_or_else(|| {
                LlmError::ProviderNotFound(format!("Client not found for provider '{selected_id}'"))
            })?;

        debug!(
            model = model_name,
            provider = %selected_id,
            strategy = ?self.strategy,
            "Routed to provider"
        );

        Ok((client, selected_id.clone()))
    }

    /// Route with failover: tries each provider until one succeeds.
    ///
    /// Returns the first successful result, or the last error if all fail.
    pub async fn route_with_failover<F, T>(
        &self,
        model_name: &str,
        f: F,
    ) -> LlmResult<T>
    where
        F: Fn(Arc<dyn LlmClient>) -> Pin<Box<dyn std::future::Future<Output = LlmResult<T>> + Send>> + Send + Sync,
        T: Send + 'static,
    {
        let provider_ids = self
            .model_providers
            .get(model_name)
            .map(|v| v.value().clone())
            .unwrap_or_default();

        if provider_ids.is_empty() {
            return Err(LlmError::ModelNotFound(format!(
                "No provider registered for model '{model_name}'"
            )));
        }

        let mut last_error = None;

        for pid in &provider_ids {
            // Skip unhealthy providers (unless recheck is due)
            if let Some(health) = self.health.get(pid) {
                if !health.is_healthy() && !health.should_recheck(self.health_cooldown_ms) {
                    continue;
                }
            }

            if let Some(client) = self.clients.get(pid).map(|c| c.value().clone()) {
                match f(client).await {
                    Ok(result) => {
                        // Mark healthy
                        if let Some(health) = self.health.get(pid) {
                            health.mark_healthy();
                        }
                        return Ok(result);
                    }
                    Err(e) => {
                        warn!(
                            provider = %pid,
                            model = model_name,
                            err = %e,
                            "Provider failed, trying next"
                        );
                        // Mark unhealthy
                        if let Some(health) = self.health.get(pid) {
                            health.mark_unhealthy();
                        }
                        last_error = Some(e);
                    }
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            LlmError::AllProvidersUnavailable(format!(
                "All providers failed for model '{model_name}'"
            ))
        }))
    }

    /// Mark a provider as healthy.
    pub fn mark_healthy(&self, provider_id: &str) {
        if let Some(h) = self.health.get(provider_id) {
            h.mark_healthy();
        }
    }

    /// Mark a provider as unhealthy.
    pub fn mark_unhealthy(&self, provider_id: &str) {
        if let Some(h) = self.health.get(provider_id) {
            h.mark_unhealthy();
        }
    }

    /// Run health checks on all registered providers.
    pub async fn check_all_health(&self) {
        let provider_ids: Vec<String> = self
            .clients
            .iter()
            .map(|e| e.key().clone())
            .collect();

        for pid in provider_ids {
            if let Some(client) = self.clients.get(&pid).map(|c| c.value().clone()) {
                match client.health_check().await {
                    Ok(true) => {
                        if let Some(h) = self.health.get(&pid) {
                            h.mark_healthy();
                        }
                        debug!(provider = %pid, "Health check passed");
                    }
                    Ok(false) => {
                        if let Some(h) = self.health.get(&pid) {
                            h.mark_unhealthy();
                        }
                        warn!(provider = %pid, "Health check failed");
                    }
                    Err(e) => {
                        if let Some(h) = self.health.get(&pid) {
                            h.mark_unhealthy();
                        }
                        warn!(provider = %pid, err = %e, "Health check error");
                    }
                }
            }
        }
    }

    /// Get health status of all providers.
    pub fn health_status(&self) -> HashMap<String, bool> {
        self.health
            .iter()
            .map(|entry| (entry.key().clone(), entry.value().is_healthy()))
            .collect()
    }

    /// Get count of registered providers.
    pub fn provider_count(&self) -> usize {
        self.clients.len()
    }

    /// Get count of registered models.
    pub fn model_count(&self) -> usize {
        self.model_providers.len()
    }

    /// Get all model names.
    pub fn model_names(&self) -> Vec<String> {
        self.model_providers
            .iter()
            .map(|e| e.key().clone())
            .collect()
    }

    /// Set routing strategy.
    pub fn set_strategy(&mut self, strategy: RoutingStrategy) {
        self.strategy = strategy;
    }

    /// Set health check cool-down period.
    pub fn set_health_cooldown_ms(&mut self, ms: u64) {
        self.health_cooldown_ms = ms;
    }
}

impl std::fmt::Debug for ProviderRouter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProviderRouter")
            .field("strategy", &self.strategy)
            .field("providers", &self.clients.len())
            .field("models", &self.model_providers.len())
            .finish()
    }
}

fn now_ms() -> u64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_routing_strategy_default() {
        assert_eq!(RoutingStrategy::default(), RoutingStrategy::Priority);
    }

    #[test]
    fn test_router_creation() {
        let router = ProviderRouter::default_router();
        assert_eq!(router.provider_count(), 0);
        assert_eq!(router.model_count(), 0);
    }

    #[test]
    fn test_provider_health() {
        let health = ProviderHealth::new();
        assert!(health.is_healthy());

        health.mark_unhealthy();
        assert!(!health.is_healthy());
        assert_eq!(health.consecutive_failures.load(Ordering::Relaxed), 1);

        health.mark_healthy();
        assert!(health.is_healthy());
        assert_eq!(health.consecutive_failures.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_route_no_providers() {
        let router = ProviderRouter::default_router();
        let result = router.route("gpt-4");
        assert!(result.is_err());
    }
}
