pe-tools 0.1.0

//! Tool trait — the interface every callable tool implements.
//!
//! The LLM sees tools via [`ToolSchema`] (name, description, JSON schema).
//! When the LLM emits a `ToolCall`, the engine looks up the tool by name
//! and calls [`Tool::execute`] with the args.
//!
//! Two implementations ship with the library:
//! - [`FunctionTool`] — wraps any async function as a tool
//! - MCP adapter (future plan) — bridges external MCP servers

use pe_core::error::PeError;
use pe_core::llm::ToolSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

/// A callable tool the LLM can invoke via tool_calls.
///
/// Implementations: [`FunctionTool`] (simple fns), MCP adapter (future),
/// graph-as-tool (Plan 014), handoff tools (Plan 009).
///
/// # Example
///
/// ```ignore
/// struct WeatherTool;
///
/// impl Tool for WeatherTool {
///     fn name(&self) -> &str { "get_weather" }
///     fn description(&self) -> &str { "Get current weather for a city" }
///     fn schema(&self) -> ToolSchema {
///         ToolSchema {
///             name: "get_weather".into(),
///             description: "Get current weather for a city".into(),
///             parameters: serde_json::json!({
///                 "type": "object",
///                 "properties": { "city": { "type": "string" } },
///                 "required": ["city"]
///             }),
///             strict: false,
///         }
///     }
///     fn execute(&self, input: Value) -> ToolFuture {
///         Box::pin(async move {
///             Ok(serde_json::json!({ "temp": 72, "unit": "F" }))
///         })
///     }
/// }
/// ```
pub trait Tool: Send + Sync {
    /// Tool name — must be unique within a [`super::ToolRegistry`].
    /// Sent to the LLM in the tools list.
    fn name(&self) -> &str;

    /// Human-readable description of what this tool does.
    /// The LLM reads this to decide when to call the tool.
    fn description(&self) -> &str;

    /// JSON Schema for this tool's input parameters.
    /// [`InjectedState`](super::InjectedState) and
    /// [`InjectedStore`](super::InjectedStore) parameters MUST be excluded.
    fn schema(&self) -> ToolSchema;

    /// Execute the tool with the given input. Called by [`super::ToolNode`].
    fn execute(&self, input: Value) -> ToolFuture;

    /// Execute with structured result metadata.
    ///
    /// Default implementation wraps `execute()` in `ToolResult::ok()`.
    /// Override to add metadata (result count, source, confidence, etc.).
    fn execute_structured(&self, input: Value) -> ToolResultFuture {
        let fut = self.execute(input);
        Box::pin(async move {
            let output = fut.await?;
            Ok(ToolResult::ok(output))
        })
    }
}

/// Type-erased async future returned by [`Tool::execute`].
pub type ToolFuture = Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>>;

/// Type-erased async future for structured tool results.
pub type ToolResultFuture = Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>;

/// Structured result from tool execution.
///
/// Wraps the tool's output value with optional metadata (result count, source,
/// confidence, execution details, etc.). Backward-compatible: `Value` auto-converts
/// to `ToolResult::ok(value)`.
///
/// # Example
///
/// ```
/// use pe_tools::tool::ToolResult;
/// use serde_json::json;
///
/// let result = ToolResult::ok(json!({"answer": 42}))
///     .with_metadata("source", json!("database"))
///     .with_metadata("result_count", json!(1));
/// assert!(result.success);
/// assert_eq!(result.metadata["source"], "database");
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    /// The tool's output value.
    pub output: Value,
    /// Whether the tool succeeded.
    pub success: bool,
    /// Tool-specific metadata (result count, source, confidence, etc.).
    #[serde(default)]
    pub metadata: HashMap<String, Value>,
}

impl ToolResult {
    /// Create a successful result.
    pub fn ok(output: Value) -> Self {
        Self {
            output,
            success: true,
            metadata: HashMap::new(),
        }
    }

    /// Create an error result with a message.
    pub fn error(msg: impl Into<String>) -> Self {
        Self {
            output: Value::String(msg.into()),
            success: false,
            metadata: HashMap::new(),
        }
    }

    /// Add metadata to this result.
    #[must_use]
    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }
}

/// Backward compatibility: raw Value auto-converts to successful ToolResult.
impl From<Value> for ToolResult {
    fn from(v: Value) -> Self {
        Self::ok(v)
    }
}

/// Type alias for the async function signature used by [`FunctionTool`].
pub type ToolFunc = Arc<
    dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>> + Send + Sync,
>;

/// Convenience wrapper: turns any async function into a [`Tool`].
///
/// # Example
///
/// ```ignore
/// let tool = FunctionTool::new(
///     "add",
///     "Add two numbers",
///     serde_json::json!({
///         "type": "object",
///         "properties": {
///             "a": { "type": "number" },
///             "b": { "type": "number" }
///         },
///         "required": ["a", "b"]
///     }),
///     |input| Box::pin(async move {
///         let a = input["a"].as_f64().unwrap_or(0.0);
///         let b = input["b"].as_f64().unwrap_or(0.0);
///         Ok(serde_json::json!(a + b))
///     }),
/// );
/// ```
pub struct FunctionTool {
    name: String,
    description: String,
    schema: ToolSchema,
    func: ToolFunc,
}

impl FunctionTool {
    /// Create a new function tool.
    ///
    /// `parameters` is the JSON Schema for the tool's input — passed to the LLM.
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: Value,
        func: impl Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, PeError>> + Send>>
        + Send
        + Sync
        + 'static,
    ) -> Self {
        let name = name.into();
        let description = description.into();
        Self {
            schema: ToolSchema {
                name: name.clone(),
                description: description.clone(),
                parameters,
                strict: false,
            },
            name,
            description,
            func: Arc::new(func),
        }
    }
}

impl Tool for FunctionTool {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        &self.description
    }

    fn schema(&self) -> ToolSchema {
        self.schema.clone()
    }

    fn execute(&self, input: Value) -> ToolFuture {
        (self.func)(input)
    }
}

/// Type alias for the async function returning a structured [`ToolResult`].
pub type StructuredToolFunc = Arc<
    dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>
        + Send
        + Sync,
>;

/// Like [`FunctionTool`], but returns structured [`ToolResult`] with metadata.
///
/// Use this when your tool needs to report result count, source, confidence, etc.
pub struct StructuredFunctionTool {
    name: String,
    description: String,
    schema: ToolSchema,
    func: StructuredToolFunc,
}

impl StructuredFunctionTool {
    /// Create a new structured function tool.
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: Value,
        func: impl Fn(Value) -> Pin<Box<dyn Future<Output = Result<ToolResult, PeError>> + Send>>
        + Send
        + Sync
        + 'static,
    ) -> Self {
        let name = name.into();
        let description = description.into();
        Self {
            schema: ToolSchema {
                name: name.clone(),
                description: description.clone(),
                parameters,
                strict: false,
            },
            name,
            description,
            func: Arc::new(func),
        }
    }
}

impl Tool for StructuredFunctionTool {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        &self.description
    }

    fn schema(&self) -> ToolSchema {
        self.schema.clone()
    }

    fn execute(&self, input: Value) -> ToolFuture {
        let func = self.func.clone();
        Box::pin(async move {
            let result = func(input).await?;
            Ok(result.output)
        })
    }

    fn execute_structured(&self, input: Value) -> ToolResultFuture {
        (self.func)(input)
    }
}

impl std::fmt::Debug for StructuredFunctionTool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StructuredFunctionTool")
            .field("name", &self.name)
            .field("description", &self.description)
            .finish()
    }
}

impl std::fmt::Debug for FunctionTool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FunctionTool")
            .field("name", &self.name)
            .field("description", &self.description)
            .finish()
    }
}

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

    #[tokio::test]
    async fn function_tool_executes_correctly() {
        let tool = FunctionTool::new(
            "add",
            "Add two numbers",
            serde_json::json!({
                "type": "object",
                "properties": {
                    "a": { "type": "number" },
                    "b": { "type": "number" }
                },
                "required": ["a", "b"]
            }),
            |input| {
                Box::pin(async move {
                    let a = input["a"].as_f64().unwrap_or(0.0);
                    let b = input["b"].as_f64().unwrap_or(0.0);
                    Ok(serde_json::json!(a + b))
                })
            },
        );

        assert_eq!(tool.name(), "add");
        assert_eq!(tool.description(), "Add two numbers");

        let schema = tool.schema();
        assert_eq!(schema.name, "add");
        assert!(!schema.strict);

        let result = tool
            .execute(serde_json::json!({"a": 3, "b": 4}))
            .await
            .unwrap();
        assert_eq!(result, serde_json::json!(7.0));
    }

    #[test]
    fn test_tool_result_ok() {
        let result = ToolResult::ok(serde_json::json!({"answer": 42}));
        assert!(result.success);
        assert_eq!(result.output, serde_json::json!({"answer": 42}));
        assert!(result.metadata.is_empty());
    }

    #[test]
    fn test_tool_result_error() {
        let result = ToolResult::error("something went wrong");
        assert!(!result.success);
        assert_eq!(
            result.output,
            serde_json::Value::String("something went wrong".into())
        );
        assert!(result.metadata.is_empty());
    }

    #[test]
    fn test_tool_result_metadata() {
        let result = ToolResult::ok(serde_json::json!({"data": [1, 2, 3]}))
            .with_metadata("source", serde_json::json!("database"))
            .with_metadata("result_count", serde_json::json!(3))
            .with_metadata("confidence", serde_json::json!(0.95));
        assert!(result.success);
        assert_eq!(result.metadata.len(), 3);
        assert_eq!(result.metadata["source"], serde_json::json!("database"));
        assert_eq!(result.metadata["result_count"], serde_json::json!(3));
        assert_eq!(result.metadata["confidence"], serde_json::json!(0.95));
    }

    #[test]
    fn test_tool_result_from_value() {
        let value = serde_json::json!({"key": "val"});
        let result: ToolResult = value.clone().into();
        assert!(result.success);
        assert_eq!(result.output, value);
        assert!(result.metadata.is_empty());
    }

    #[tokio::test]
    async fn test_execute_structured_default() {
        let tool = FunctionTool::new(
            "add",
            "Add two numbers",
            serde_json::json!({"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}),
            |input| {
                Box::pin(async move {
                    let a = input["a"].as_f64().unwrap_or(0.0);
                    let b = input["b"].as_f64().unwrap_or(0.0);
                    Ok(serde_json::json!(a + b))
                })
            },
        );

        let result = tool
            .execute_structured(serde_json::json!({"a": 3, "b": 4}))
            .await
            .unwrap();
        assert!(result.success);
        assert_eq!(result.output, serde_json::json!(7.0));
        assert!(result.metadata.is_empty());
    }

    #[tokio::test]
    async fn function_tool_propagates_error() {
        let tool = FunctionTool::new(
            "fail",
            "Always fails",
            serde_json::json!({"type": "object"}),
            |_input| {
                Box::pin(async move {
                    Err(PeError::ToolExecution {
                        tool: "fail".into(),
                        reason: "intentional failure".into(),
                    })
                })
            },
        );

        let result = tool.execute(serde_json::json!({})).await;
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(err.to_string().contains("intentional failure"));
    }
}