use async_trait::async_trait;
use std::sync::Arc;
use tracing::{info, warn};

use swarm_core::error::SwarmResult;

use crate::provider::{CompletionRequest, CompletionResponse, LlmProvider};

/// A provider that tries multiple backends in order until one succeeds.
pub struct FallbackChain {
    name: String,
    providers: Vec<ChainEntry>,
}

struct ChainEntry {
    label: String,
    model: String,
    provider: Arc<dyn LlmProvider>,
}

impl FallbackChain {
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            providers: Vec::new(),
        }
    }

    pub fn add(mut self, label: impl Into<String>, model: impl Into<String>, provider: Arc<dyn LlmProvider>) -> Self {
        self.providers.push(ChainEntry {
            label: label.into(),
            model: model.into(),
            provider,
        });
        self
    }

    pub fn is_empty(&self) -> bool {
        self.providers.is_empty()
    }
}

#[async_trait]
impl LlmProvider for FallbackChain {
    fn name(&self) -> &str {
        &self.name
    }

    async fn complete(&self, request: CompletionRequest) -> SwarmResult<CompletionResponse> {
        let mut last_error = None;

        for (i, entry) in self.providers.iter().enumerate() {
            // Override the model in the request with this entry's model
            let req = CompletionRequest {
                model: entry.model.clone(),
                messages: request.messages.clone(),
                system: request.system.clone(),
                tools: request.tools.clone(),
                temperature: request.temperature,
                max_tokens: request.max_tokens,
            };

            match entry.provider.complete(req).await {
                Ok(response) => {
                    if i > 0 {
                        info!(
                            chain = %self.name,
                            provider = %entry.label,
                            "Fallback succeeded on attempt {}",
                            i + 1
                        );
                    }
                    return Ok(response);
                }
                Err(e) => {
                    warn!(
                        chain = %self.name,
                        provider = %entry.label,
                        error = %e,
                        "Provider failed, trying next"
                    );
                    last_error = Some(e);
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            swarm_core::error::SwarmError::LlmError("No providers in fallback chain".into())
        }))
    }
}
