use async_trait::async_trait;
use std::sync::Arc;

use crate::product::agent::function_tool::FunctionCallError;
use crate::product::agent::mcp_tool_call::handle_mcp_tool_call;
use crate::product::agent::tools::context::ToolInvocation;
use crate::product::agent::tools::context::ToolOutput;
use crate::product::agent::tools::context::ToolPayload;
use crate::product::agent::tools::registry::ToolHandler;
use crate::product::agent::tools::registry::ToolKind;
use lha_llm::ToolResultPayload;

pub struct McpHandler;

#[async_trait]
impl ToolHandler for McpHandler {
    fn kind(&self) -> ToolKind {
        ToolKind::Mcp
    }

    async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
        let ToolInvocation {
            session,
            turn,
            call_id,
            payload,
            ..
        } = invocation;

        let payload = match payload {
            ToolPayload::Mcp {
                server,
                tool,
                raw_arguments,
            } => (server, tool, raw_arguments),
            _ => {
                return Err(FunctionCallError::RespondToModel(
                    "mcp handler received unsupported payload".to_string(),
                ));
            }
        };

        let (server, tool, raw_arguments) = payload;
        let arguments_str = raw_arguments;

        let response = handle_mcp_tool_call(
            Arc::clone(&session),
            turn.as_ref(),
            call_id.clone(),
            server,
            tool,
            arguments_str,
        )
        .await;

        match response.payload {
            ToolResultPayload::Structured {
                content,
                content_items,
                success,
            } => Ok(ToolOutput::Function {
                content,
                content_items,
                success,
            }),
            ToolResultPayload::Text { output } => Ok(ToolOutput::Function {
                content: output,
                content_items: None,
                success: Some(true),
            }),
        }
    }
}
