//! LLM Provider implementations.
//!
//! Provides two LLM providers:
//!
//! 1. **`SamplingLlmProvider`** — Uses MCP sampling (`sampling/createMessage`)
//!    to request the host client (Claude, GPT, etc.) to generate text. This is
//!    the primary provider that gives agents real LLM completions.
//!
//! 2. **`PassthroughLlmProvider`** — Fallback provider that signals LLM is
//!    unavailable. `is_available()` returns `false` and all completions return
//!    errors. Used when the host client doesn't support MCP sampling.
//!
//! **Failure model**: If the host returns an error on the first
//! `sampling/createMessage` call, it means the host doesn't support sampling.
//! `SamplingLlmProvider` permanently disables itself (`is_available() = false`)
//! and all agents degrade to their heuristic-only paths.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering::Relaxed};

use anyhow::Result;
use rmcp::model::{
    Content, CreateMessageRequestParam, RawContent, Role, SamplingMessage,
};
use rmcp::service::Peer;
use rmcp::RoleServer;
use tokio::sync::OnceCell;
use tracing::{debug, info, warn};

use crate::types::LlmProvider;

// ---------------------------------------------------------------------------
// SamplingLlmProvider
// ---------------------------------------------------------------------------

/// LLM provider that uses MCP sampling to request completions from the host.
///
/// The host client (e.g. Claude Code, Claude Desktop) receives a
/// `sampling/createMessage` request and generates text using its own LLM.
/// This gives FURNACE agents (Trinity, Gimli, SearchOrchestrator, etc.)
/// access to real LLM completions without needing direct API keys.
///
/// **One-shot failure**: If sampling fails for any reason (peer not connected,
/// host returns error), the provider permanently disables itself. A failure
/// means the host doesn't support `sampling/createMessage` — retrying would
/// just waste time. All agents check `is_available()` and fall back to
/// heuristic-only paths.
pub struct SamplingLlmProvider {
    peer: Arc<OnceCell<Peer<RoleServer>>>,
    /// Set to `true` permanently on first successful sampling call.
    sampling_confirmed: AtomicBool,
    /// Set to `true` permanently on first failure — host doesn't support sampling.
    sampling_failed: AtomicBool,
}

impl SamplingLlmProvider {
    /// Create a new SamplingLlmProvider with a shared peer cell.
    ///
    /// The peer cell is populated by `FurnaceHandler::call_tool` on the first
    /// tool call, when the rmcp framework provides the peer via RequestContext.
    pub fn new(peer: Arc<OnceCell<Peer<RoleServer>>>) -> Self {
        Self {
            peer,
            sampling_confirmed: AtomicBool::new(false),
            sampling_failed: AtomicBool::new(false),
        }
    }

    /// Extract text content from an MCP Content value.
    fn extract_text(content: &Content) -> Option<String> {
        if let RawContent::Text(t) = &content.raw {
            Some(t.text.clone())
        } else {
            None
        }
    }
}

#[async_trait::async_trait]
impl LlmProvider for SamplingLlmProvider {
    fn name(&self) -> &str {
        match std::env::var("FURNACE_LLM_PROVIDER") {
            Ok(ref v) if !v.is_empty() => {
                Box::leak(v.clone().into_boxed_str())
            }
            _ => "claude",
        }
    }

    fn is_available(&self) -> bool {
        // Once failed, permanently unavailable — host doesn't support sampling.
        if self.sampling_failed.load(Relaxed) {
            return false;
        }
        // Once confirmed working, always available.
        if self.sampling_confirmed.load(Relaxed) {
            return true;
        }
        // Not yet tested — available if peer is connected (first call will decide).
        self.peer.get().is_some()
    }

    async fn complete(&self, prompt: &str) -> Result<String> {
        self.complete_with_system("", prompt).await
    }

    async fn complete_with_system(&self, system: &str, prompt: &str) -> Result<String> {
        // Fast path: already known to be unsupported.
        if self.sampling_failed.load(Relaxed) {
            return Err(anyhow::anyhow!(
                "MCP sampling permanently disabled: host does not support sampling/createMessage"
            ));
        }

        let peer = match self.peer.get() {
            Some(p) => p.clone(),
            None => {
                self.sampling_failed.store(true, Relaxed);
                warn!("MCP sampling unavailable: peer not connected — disabling permanently");
                return Err(anyhow::anyhow!("MCP sampling unavailable: peer not connected"));
            }
        };

        let params = CreateMessageRequestParam {
            messages: vec![SamplingMessage {
                role: Role::User,
                content: Content::text(prompt),
            }],
            model_preferences: None,
            system_prompt: if system.is_empty() {
                None
            } else {
                Some(system.to_string())
            },
            include_context: Some("thisServer".to_string()),
            temperature: Some(0.2),
            max_tokens: 16384,
            stop_sequences: None,
            metadata: Some(serde_json::json!({
                "furnace_autonomous": true,
                "pipeline": "autoforge",
            })),
        };

        debug!(
            system_len = system.len(),
            prompt_len = prompt.len(),
            "Sending MCP sampling request"
        );

        match peer.create_message(params).await {
            Ok(result) => {
                if !self.sampling_confirmed.load(Relaxed) {
                    self.sampling_confirmed.store(true, Relaxed);
                    info!("MCP sampling confirmed working — LLM features enabled");
                }
                let text = Self::extract_text(&result.message.content)
                    .unwrap_or_default();
                debug!(
                    response_len = text.len(),
                    model = %result.model,
                    "MCP sampling response received"
                );
                Ok(text)
            }
            Err(e) => {
                self.sampling_failed.store(true, Relaxed);
                warn!(
                    error = %e,
                    "MCP sampling failed — host does not support sampling/createMessage. \
                     Disabling LLM permanently; agents will use heuristic-only paths."
                );
                Err(anyhow::anyhow!("MCP sampling failed: {e}"))
            }
        }
    }
}

// ---------------------------------------------------------------------------
// PassthroughLlmProvider
// ---------------------------------------------------------------------------

/// Fallback LLM provider that signals LLM is unavailable.
///
/// `is_available()` always returns `false` and all completion calls return
/// errors. This is the effective state that `SamplingLlmProvider` degrades
/// to after a sampling failure — agents check `is_available()` and fall
/// back to heuristic-only code paths instead of calling `complete()`.
pub struct PassthroughLlmProvider;

#[async_trait::async_trait]
impl LlmProvider for PassthroughLlmProvider {
    fn name(&self) -> &str {
        match std::env::var("FURNACE_LLM_PROVIDER") {
            Ok(ref v) if !v.is_empty() => {
                Box::leak(v.clone().into_boxed_str())
            }
            _ => "claude",
        }
    }

    fn is_available(&self) -> bool {
        false
    }

    async fn complete(&self, _prompt: &str) -> Result<String> {
        Err(anyhow::anyhow!(
            "LLM unavailable: MCP sampling not supported by host client. \
             The host client must support sampling/createMessage."
        ))
    }

    async fn complete_with_system(&self, _system: &str, _prompt: &str) -> Result<String> {
        Err(anyhow::anyhow!(
            "LLM unavailable: MCP sampling not supported by host client. \
             The host client must support sampling/createMessage."
        ))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::Ordering::Relaxed;

    // -- SamplingLlmProvider state machine tests --

    #[test]
    fn test_sampling_provider_initial_state_no_peer() {
        let peer = Arc::new(OnceCell::new());
        let provider = SamplingLlmProvider::new(Arc::clone(&peer));

        assert!(!provider.sampling_confirmed.load(Relaxed));
        assert!(!provider.sampling_failed.load(Relaxed));
        assert!(!provider.is_available(), "No peer → not available");
    }

    #[test]
    fn test_sampling_provider_permanently_disabled_after_failure() {
        let peer = Arc::new(OnceCell::new());
        let provider = SamplingLlmProvider::new(Arc::clone(&peer));

        provider.sampling_failed.store(true, Relaxed);

        assert!(!provider.is_available(), "After failure → permanently unavailable");
    }

    #[test]
    fn test_sampling_provider_available_after_success() {
        let peer = Arc::new(OnceCell::new());
        let provider = SamplingLlmProvider::new(Arc::clone(&peer));

        provider.sampling_confirmed.store(true, Relaxed);

        assert!(provider.is_available(), "After confirmed success → always available");
    }

    #[test]
    fn test_sampling_failure_overrides_confirmed() {
        let peer = Arc::new(OnceCell::new());
        let provider = SamplingLlmProvider::new(Arc::clone(&peer));

        // If somehow both are set, failure wins (checked first).
        provider.sampling_confirmed.store(true, Relaxed);
        provider.sampling_failed.store(true, Relaxed);

        assert!(!provider.is_available(), "Failed flag takes precedence");
    }

    #[tokio::test]
    async fn test_sampling_provider_disables_permanently_on_no_peer() {
        let peer = Arc::new(OnceCell::new());
        let provider = SamplingLlmProvider::new(Arc::clone(&peer));

        let result = provider.complete_with_system("sys", "prompt").await;
        assert!(result.is_err());
        assert!(provider.sampling_failed.load(Relaxed), "Should be permanently disabled");
        assert!(!provider.sampling_confirmed.load(Relaxed), "Should not be confirmed");

        // Subsequent calls fail immediately without trying.
        let result2 = provider.complete("anything").await;
        assert!(result2.is_err());
        assert!(result2.unwrap_err().to_string().contains("permanently disabled"));
    }

    // -- PassthroughLlmProvider tests --

    #[test]
    fn test_passthrough_provider_is_never_available() {
        let provider = PassthroughLlmProvider;
        assert!(!provider.is_available());
    }

    #[tokio::test]
    async fn test_passthrough_provider_complete_returns_error() {
        let provider = PassthroughLlmProvider;
        let result = provider.complete("hello").await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("MCP sampling not supported"));
    }

    #[tokio::test]
    async fn test_passthrough_provider_complete_with_system_returns_error() {
        let provider = PassthroughLlmProvider;
        let result = provider.complete_with_system("system", "hello").await;
        assert!(result.is_err());
    }

    // -- LlmProvider trait default tests --

    #[test]
    fn test_llm_provider_trait_default_is_available() {
        use crate::types::LlmProvider;

        struct DefaultMock;

        #[async_trait::async_trait]
        impl LlmProvider for DefaultMock {
            async fn complete(&self, _: &str) -> Result<String> { Ok("ok".into()) }
            async fn complete_with_system(&self, _: &str, _: &str) -> Result<String> { Ok("ok".into()) }
        }

        let mock = DefaultMock;
        assert!(mock.is_available(), "Default trait impl should return true");
    }
}
