prvng_platform/crates/rag/src/tools.rs

//! Claude function calling / tool use support for RAG system
//!
//! Provides a framework for Claude to invoke system operations safely with:
//! - Tool trait for extensible tool implementations
//! - Tool registry for managing available tools
//! - Security validation and rate limiting
//! - Audit logging for all tool invocations
//! - Result aggregation and error handling

use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{info, warn};

use crate::error::{RagError, Result};

/// Tool input/output types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInput {
    /// Tool parameters as JSON
    pub params: serde_json::Value,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolOutput {
    /// Tool execution result
    pub result: String,
    /// Whether execution was successful
    pub success: bool,
    /// Optional error message
    pub error: Option<String>,
}

/// Tool definition for Claude to understand
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name (used by Claude)
    pub name: String,
    /// Tool description for Claude
    pub description: String,
    /// JSON schema for input parameters
    pub input_schema: serde_json::Value,
}

/// Tool implementation trait
#[async_trait]
pub trait RagTool: Send + Sync {
    /// Get tool definition (name, description, schema)
    fn definition(&self) -> ToolDefinition;

    /// Execute the tool with given input
    async fn execute(&self, input: ToolInput) -> Result<ToolOutput>;

    /// Validate if user is authorized to call this tool
    async fn validate_access(&self, user_id: &str) -> Result<bool> {
        // Default: allow all authenticated users
        Ok(!user_id.is_empty())
    }

    /// Check if tool is currently rate-limited
    async fn check_rate_limit(&self, _user_id: &str) -> Result<bool> {
        // Default: no rate limiting
        Ok(true)
    }

    /// Record audit log for tool invocation
    async fn audit_log(&self, user_id: &str, _input: &ToolInput, output: &ToolOutput) {
        info!(
            target: "rag::tools::audit",
            tool = %self.definition().name,
            user = user_id,
            success = output.success,
            "Tool execution logged"
        );
    }
}

/// Tool registry for managing available tools
pub struct ToolRegistry {
    tools: Arc<RwLock<HashMap<String, Arc<dyn RagTool>>>>,
    call_stats: Arc<RwLock<ToolCallStats>>,
}

/// Statistics for tool calls
#[derive(Debug, Clone, Default)]
pub struct ToolCallStats {
    /// Total tool calls
    pub total_calls: u64,
    /// Successful tool calls
    pub successful_calls: u64,
    /// Failed tool calls
    pub failed_calls: u64,
    /// Blocked calls (auth/rate limit)
    pub blocked_calls: u64,
    /// Calls by tool name
    pub calls_by_tool: HashMap<String, u64>,
}

impl ToolRegistry {
    /// Create a new tool registry
    pub fn new() -> Self {
        Self {
            tools: Arc::new(RwLock::new(HashMap::new())),
            call_stats: Arc::new(RwLock::new(ToolCallStats::default())),
        }
    }

    /// Register a new tool
    pub async fn register(&self, tool: Arc<dyn RagTool>) -> Result<()> {
        let def = tool.definition();
        let mut tools = self.tools.write().await;

        if tools.contains_key(&def.name) {
            return Err(RagError::ToolError(format!(
                "Tool '{}' already registered",
                def.name
            )));
        }

        tools.insert(def.name.clone(), tool);
        info!("Tool registered: {}", def.name);
        Ok(())
    }

    /// Get tool by name
    pub async fn get_tool(&self, name: &str) -> Result<Arc<dyn RagTool>> {
        let tools = self.tools.read().await;
        tools
            .get(name)
            .cloned()
            .ok_or_else(|| RagError::ToolError(format!("Tool '{}' not found", name)))
    }

    /// Get all tool definitions
    pub async fn get_definitions(&self) -> Vec<ToolDefinition> {
        let tools = self.tools.read().await;
        tools.values().map(|t| t.definition()).collect()
    }

    /// Call a tool with security validation
    pub async fn call_tool(
        &self,
        name: &str,
        input: ToolInput,
        user_id: &str,
    ) -> Result<ToolOutput> {
        let tool = self.get_tool(name).await?;

        // 1. Validate access
        if !tool.validate_access(user_id).await? {
            warn!(user = user_id, tool = name, "Tool access denied");
            let mut stats = self.call_stats.write().await;
            stats.blocked_calls += 1;
            return Ok(ToolOutput {
                result: "Access denied".to_string(),
                success: false,
                error: Some("User not authorized for this tool".to_string()),
            });
        }

        // 2. Check rate limiting
        if !tool.check_rate_limit(user_id).await? {
            warn!(user = user_id, tool = name, "Tool rate limited");
            let mut stats = self.call_stats.write().await;
            stats.blocked_calls += 1;
            return Ok(ToolOutput {
                result: "Rate limited".to_string(),
                success: false,
                error: Some("Tool call rate limit exceeded".to_string()),
            });
        }

        // 3. Execute tool
        let output = tool.execute(input.clone()).await?;

        // 4. Record audit log
        tool.audit_log(user_id, &input, &output).await;

        // 5. Update statistics
        let mut stats = self.call_stats.write().await;
        stats.total_calls += 1;
        *stats.calls_by_tool.entry(name.to_string()).or_insert(0) += 1;
        if output.success {
            stats.successful_calls += 1;
        } else {
            stats.failed_calls += 1;
        }

        Ok(output)
    }

    /// Get tool call statistics
    pub async fn get_stats(&self) -> ToolCallStats {
        self.call_stats.read().await.clone()
    }

    /// Reset statistics
    pub async fn reset_stats(&self) {
        let mut stats = self.call_stats.write().await;
        *stats = ToolCallStats::default();
    }
}

impl Default for ToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// CORE TOOL IMPLEMENTATIONS
// ============================================================================

/// Create server tool - allows Claude to provision infrastructure
pub struct CreateServerTool {
    description: String,
}

impl CreateServerTool {
    pub fn new() -> Self {
        Self {
            description: "Provision a new server with specified configuration".to_string(),
        }
    }
}

impl Default for CreateServerTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl RagTool for CreateServerTool {
    fn definition(&self) -> ToolDefinition {
        ToolDefinition {
            name: "create_server".to_string(),
            description: "Create a new server with specified CPU, memory, and location".to_string(),
            input_schema: serde_json::json!({
                "type": "object",
                "properties": {
                    "hostname": {
                        "type": "string",
                        "description": "Server hostname"
                    },
                    "cores": {
                        "type": "integer",
                        "description": "Number of CPU cores (1-128)",
                        "minimum": 1,
                        "maximum": 128
                    },
                    "memory_gb": {
                        "type": "integer",
                        "description": "Memory in GB (1-2048)",
                        "minimum": 1,
                        "maximum": 2048
                    },
                    "region": {
                        "type": "string",
                        "description": "Server region/zone"
                    }
                },
                "required": ["hostname", "cores", "memory_gb", "region"]
            }),
        }
    }

    async fn execute(&self, input: ToolInput) -> Result<ToolOutput> {
        // Extract parameters
        let hostname = input
            .params
            .get("hostname")
            .and_then(|v| v.as_str())
            .ok_or_else(|| RagError::ToolError("Missing hostname".to_string()))?;

        let cores = input
            .params
            .get("cores")
            .and_then(|v| v.as_i64())
            .ok_or_else(|| RagError::ToolError("Missing cores".to_string()))?
            as i32;

        let memory_gb = input
            .params
            .get("memory_gb")
            .and_then(|v| v.as_i64())
            .ok_or_else(|| RagError::ToolError("Missing memory_gb".to_string()))?
            as i32;

        let region = input
            .params
            .get("region")
            .and_then(|v| v.as_str())
            .ok_or_else(|| RagError::ToolError("Missing region".to_string()))?;

        // Validate parameters
        if !(1..=128).contains(&cores) {
            return Ok(ToolOutput {
                result: format!("Invalid cores: {}", cores),
                success: false,
                error: Some("Cores must be between 1 and 128".to_string()),
            });
        }

        if !(1..=2048).contains(&memory_gb) {
            return Ok(ToolOutput {
                result: format!("Invalid memory: {}GB", memory_gb),
                success: false,
                error: Some("Memory must be between 1 and 2048 GB".to_string()),
            });
        }

        // In production, this would call the provisioning API
        // For now, return simulation
        let result = format!(
            "Server '{}' created: {} cores, {}GB memory in region {}",
            hostname, cores, memory_gb, region
        );

        info!("Server creation simulated: {}", result);

        Ok(ToolOutput {
            result,
            success: true,
            error: None,
        })
    }

    async fn validate_access(&self, user_id: &str) -> Result<bool> {
        // Only admin users can create servers
        Ok(user_id.ends_with("@admin"))
    }
}

/// Workspace status tool - allows Claude to check workspace info
pub struct WorkspaceStatusTool {
    workspace_name: String,
}

impl WorkspaceStatusTool {
    pub fn new(workspace_name: String) -> Self {
        Self { workspace_name }
    }
}

#[async_trait]
impl RagTool for WorkspaceStatusTool {
    fn definition(&self) -> ToolDefinition {
        ToolDefinition {
            name: "get_workspace_status".to_string(),
            description: "Get current workspace status and metrics".to_string(),
            input_schema: serde_json::json!({
                "type": "object",
                "properties": {
                    "include_metrics": {
                        "type": "boolean",
                        "description": "Include detailed metrics"
                    }
                }
            }),
        }
    }

    async fn execute(&self, input: ToolInput) -> Result<ToolOutput> {
        let include_metrics = input
            .params
            .get("include_metrics")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let status = if include_metrics {
            format!(
                "Workspace '{}': 5 servers, 3 taskservs, 2 clusters. Metrics: 99.8% uptime",
                self.workspace_name
            )
        } else {
            format!(
                "Workspace '{}': 5 servers, 3 taskservs, 2 clusters",
                self.workspace_name
            )
        };

        info!("Workspace status retrieved: {}", status);

        Ok(ToolOutput {
            result: status,
            success: true,
            error: None,
        })
    }
}

/// Taskserv management tool - allows Claude to manage infrastructure services
pub struct TaskservManagementTool {
    operations: Vec<String>,
}

impl TaskservManagementTool {
    pub fn new() -> Self {
        Self {
            operations: vec![
                "create".to_string(),
                "delete".to_string(),
                "update".to_string(),
                "restart".to_string(),
            ],
        }
    }
}

impl Default for TaskservManagementTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl RagTool for TaskservManagementTool {
    fn definition(&self) -> ToolDefinition {
        ToolDefinition {
            name: "manage_taskserv".to_string(),
            description: "Manage infrastructure task services (create, delete, update, restart)"
                .to_string(),
            input_schema: serde_json::json!({
                "type": "object",
                "properties": {
                    "operation": {
                        "type": "string",
                        "enum": ["create", "delete", "update", "restart"],
                        "description": "Operation to perform"
                    },
                    "service": {
                        "type": "string",
                        "description": "Service name (e.g., kubernetes, containerd)"
                    },
                    "version": {
                        "type": "string",
                        "description": "Service version (optional)"
                    }
                },
                "required": ["operation", "service"]
            }),
        }
    }

    async fn execute(&self, input: ToolInput) -> Result<ToolOutput> {
        let operation = input
            .params
            .get("operation")
            .and_then(|v| v.as_str())
            .ok_or_else(|| RagError::ToolError("Missing operation".to_string()))?;

        let service = input
            .params
            .get("service")
            .and_then(|v| v.as_str())
            .ok_or_else(|| RagError::ToolError("Missing service".to_string()))?;

        // Validate operation
        if !self.operations.contains(&operation.to_string()) {
            return Ok(ToolOutput {
                result: format!("Invalid operation: {}", operation),
                success: false,
                error: Some(format!(
                    "Operation must be one of: {}",
                    self.operations.join(", ")
                )),
            });
        }

        let result = format!("Taskserv '{}' operation '{}' completed", service, operation);

        info!("Taskserv operation simulated: {}", result);

        Ok(ToolOutput {
            result,
            success: true,
            error: None,
        })
    }

    async fn validate_access(&self, _user_id: &str) -> Result<bool> {
        // Delete operation requires admin
        Ok(true) // Simplified for now
    }
}

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

    #[test]
    fn test_tool_input_creation() {
        let input = ToolInput {
            params: serde_json::json!({
                "hostname": "server-01",
                "cores": 4
            }),
        };

        assert_eq!(input.params["hostname"], "server-01");
        assert_eq!(input.params["cores"], 4);
    }

    #[test]
    fn test_tool_output_success() {
        let output = ToolOutput {
            result: "Operation successful".to_string(),
            success: true,
            error: None,
        };

        assert!(output.success);
        assert_eq!(output.result, "Operation successful");
        assert!(output.error.is_none());
    }

    #[test]
    fn test_tool_output_failure() {
        let output = ToolOutput {
            result: "Operation failed".to_string(),
            success: false,
            error: Some("Invalid parameters".to_string()),
        };

        assert!(!output.success);
        assert!(output.error.is_some());
    }

    #[tokio::test]
    async fn test_registry_creation() {
        let registry = ToolRegistry::new();
        let definitions = registry.get_definitions().await;
        assert_eq!(definitions.len(), 0);
    }

    #[tokio::test]
    async fn test_tool_registration() {
        let registry = ToolRegistry::new();
        let tool = Arc::new(CreateServerTool::new());

        registry.register(tool).await.unwrap();
        let definitions = registry.get_definitions().await;
        assert_eq!(definitions.len(), 1);
        assert_eq!(definitions[0].name, "create_server");
    }

    #[tokio::test]
    async fn test_duplicate_registration() {
        let registry = ToolRegistry::new();
        let tool: Arc<dyn RagTool> = Arc::new(CreateServerTool::new());

        registry.register(Arc::clone(&tool)).await.unwrap();
        let result = registry.register(Arc::clone(&tool)).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_create_server_tool() {
        let tool = CreateServerTool::new();
        let def = tool.definition();

        assert_eq!(def.name, "create_server");
        assert!(def.description.contains("server"));
    }

    #[tokio::test]
    async fn test_create_server_valid_input() {
        let tool = CreateServerTool::new();
        let input = ToolInput {
            params: serde_json::json!({
                "hostname": "web-01",
                "cores": 4,
                "memory_gb": 8,
                "region": "us-east-1"
            }),
        };

        let output = tool.execute(input).await.unwrap();
        assert!(output.success);
        assert!(output.result.contains("web-01"));
    }

    #[tokio::test]
    async fn test_create_server_invalid_cores() {
        let tool = CreateServerTool::new();
        let input = ToolInput {
            params: serde_json::json!({
                "hostname": "web-01",
                "cores": 256,
                "memory_gb": 8,
                "region": "us-east-1"
            }),
        };

        let output = tool.execute(input).await.unwrap();
        assert!(!output.success);
        assert!(output.error.is_some());
    }

    #[tokio::test]
    async fn test_workspace_status_tool() {
        let tool = WorkspaceStatusTool::new("production".to_string());
        let input = ToolInput {
            params: serde_json::json!({"include_metrics": true}),
        };

        let output = tool.execute(input).await.unwrap();
        assert!(output.success);
        assert!(output.result.contains("production"));
    }

    #[tokio::test]
    async fn test_taskserv_management_tool() {
        let tool = TaskservManagementTool::new();
        let input = ToolInput {
            params: serde_json::json!({
                "operation": "create",
                "service": "kubernetes",
                "version": "1.28.0"
            }),
        };

        let output = tool.execute(input).await.unwrap();
        assert!(output.success);
        assert!(output.result.contains("kubernetes"));
    }

    #[tokio::test]
    async fn test_taskserv_invalid_operation() {
        let tool = TaskservManagementTool::new();
        let input = ToolInput {
            params: serde_json::json!({
                "operation": "invalid",
                "service": "kubernetes"
            }),
        };

        let output = tool.execute(input).await.unwrap();
        assert!(!output.success);
    }

    #[tokio::test]
    async fn test_tool_registry_call() {
        let registry = ToolRegistry::new();
        let tool = Arc::new(WorkspaceStatusTool::new("test".to_string()));
        registry.register(tool).await.unwrap();

        let input = ToolInput {
            params: serde_json::json!({"include_metrics": false}),
        };

        let output = registry
            .call_tool("get_workspace_status", input, "user@example.com")
            .await
            .unwrap();

        assert!(output.success);
    }

    #[tokio::test]
    async fn test_create_server_access_control() {
        let tool = CreateServerTool::new();

        // Admin user should have access (user_id ends with @admin)
        let admin_access = tool.validate_access("admin@admin").await.unwrap();
        assert!(admin_access);

        // Non-admin should not have access
        let user_access = tool.validate_access("user@example.com").await.unwrap();
        assert!(!user_access);
    }

    #[tokio::test]
    async fn test_tool_statistics() {
        let registry = ToolRegistry::new();
        let tool = Arc::new(WorkspaceStatusTool::new("test".to_string()));
        registry.register(tool).await.unwrap();

        let input = ToolInput {
            params: serde_json::json!({"include_metrics": false}),
        };

        registry
            .call_tool("get_workspace_status", input.clone(), "user@test.com")
            .await
            .unwrap();

        let stats = registry.get_stats().await;
        assert_eq!(stats.total_calls, 1);
        assert_eq!(stats.successful_calls, 1);
        assert_eq!(stats.failed_calls, 0);
    }

    #[test]
    fn test_tool_call_stats_default() {
        let stats = ToolCallStats::default();
        assert_eq!(stats.total_calls, 0);
        assert_eq!(stats.successful_calls, 0);
        assert_eq!(stats.failed_calls, 0);
        assert_eq!(stats.blocked_calls, 0);
    }
}