use crate::{
    LLMProvider, Tool,
    chat::{ChatMessage, ChatProvider, ChatResponse, Content, StreamChunk},
    completion::{CompletionProvider, CompletionRequest, CompletionResponse},
    embedding::EmbeddingProvider,
    error::LLMError,
    stt, tts,
};
use anyhow::Result;
use async_trait::async_trait;
use futures::Stream;
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;

/// Adapter interface for your host‐side implementations.
///
/// Tool implementations return `Vec<Content>` to support mixed content results
/// (e.g., text + images from MCP tools).
#[async_trait]
pub trait CallFunctionTool: Send + Sync {
    fn descriptor(&self) -> Tool;
    async fn call(&self, args: Value) -> anyhow::Result<Vec<Content>>;

    /// Returns the server name for server-aware tools (e.g., MCP tools).
    /// Returns None for tools that don't have server information.
    fn server_name(&self) -> Option<&str> {
        None
    }
}

/*
#[derive(Clone, Serialize, JsonSchema)]
pub struct CallableFunctionTool {
    /// Flattened into the JSON so the LLM sees `name`, `description`, `parameters`
    #[serde(flatten)]
    pub meta: FunctionTool,

    /// Not serialized; the actual Rust callback logic
    #[serde(skip)]
    callback:
        Arc<dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> + Send + Sync>,
}

impl CallableFunctionTool {
    /// Create a new one by giving it its metadata **and** an async function
    pub fn new<F, Fut>(meta: FunctionTool, f: F) -> Self
    where
        F: Fn(Value) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<String>> + Send + 'static,
    {
        Self {
            meta,
            callback: Arc::new(move |args| Box::pin(f(args))),
        }
    }
}

#[async_trait]
impl CallFunctionTool for CallableFunctionTool {
    fn descriptor(&self) -> FunctionTool {
        self.meta.clone()
    }

    async fn call(&self, args: Value) -> Result<String> {
        (self.callback)(args).await
    }
}
*/

pub struct ToolEnabledProvider {
    inner: Box<dyn LLMProvider + Send + Sync>,
    registry: HashMap<String, Box<dyn CallFunctionTool>>,
    tool_list: Vec<Tool>,
}

impl ToolEnabledProvider {
    pub fn new(
        inner: Box<dyn LLMProvider + Send + Sync>,
        registry: HashMap<String, Box<dyn CallFunctionTool>>,
    ) -> Self {
        let tool_list = registry.values().map(|t| t.descriptor()).collect();
        ToolEnabledProvider {
            inner,
            registry,
            tool_list,
        }
    }

    /// Construct with a pre-built tool list.
    pub fn with_tool_list(
        inner: Box<dyn LLMProvider + Send + Sync>,
        registry: HashMap<String, Box<dyn CallFunctionTool>>,
        tool_list: Vec<Tool>,
    ) -> Self {
        ToolEnabledProvider {
            inner,
            registry,
            tool_list,
        }
    }
}

#[async_trait]
impl CompletionProvider for ToolEnabledProvider {
    async fn complete(&self, req: &CompletionRequest) -> Result<CompletionResponse, LLMError> {
        self.inner.complete(req).await
    }
}

#[async_trait]
impl EmbeddingProvider for ToolEnabledProvider {
    async fn embed(&self, input: Vec<String>) -> Result<Vec<Vec<f32>>, LLMError> {
        self.inner.embed(input).await
    }
}

#[async_trait]
impl ChatProvider for ToolEnabledProvider {
    fn supports_streaming(&self) -> bool {
        self.inner.supports_streaming()
    }

    async fn chat_with_tools(
        &self,
        messages: &[ChatMessage],
        tools: Option<&[Tool]>,
    ) -> Result<Box<dyn ChatResponse>, LLMError> {
        let to_send = tools.unwrap_or(&self.tool_list);
        self.inner.chat_with_tools(messages, Some(to_send)).await
    }

    async fn chat_stream_with_tools(
        &self,
        messages: &[ChatMessage],
        tools: Option<&[Tool]>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send>>, LLMError> {
        let to_send = tools.unwrap_or(&self.tool_list);
        self.inner
            .chat_stream_with_tools(messages, Some(to_send))
            .await
    }
}

#[async_trait]
impl LLMProvider for ToolEnabledProvider {
    fn tools(&self) -> Option<&[Tool]> {
        Some(&self.tool_list)
    }

    async fn call_tool(
        &self,
        name: &str,
        args: serde_json::Value,
    ) -> Result<Vec<Content>, LLMError> {
        let tool = self
            .registry
            .get(name)
            .ok_or_else(|| LLMError::InvalidRequest(format!("unknown tool `{}`", name)))?;
        tool.call(args)
            .await
            .map_err(|e| LLMError::InvalidRequest(e.to_string()))
    }

    fn tool_server_name(&self, name: &str) -> Option<&str> {
        self.registry.get(name).and_then(|tool| tool.server_name())
    }

    async fn transcribe(&self, req: &stt::SttRequest) -> Result<stt::SttResponse, LLMError> {
        self.inner.transcribe(req).await
    }

    async fn speech(&self, req: &tts::TtsRequest) -> Result<tts::TtsResponse, LLMError> {
        self.inner.speech(req).await
    }
}
