prvng_platform/crates/ai-service/src/mcp.rs

//! MCP (Model Context Protocol) tool registry and execution
//!
//! Provides tool definition, registration, and execution for RAG, Guidance,
//! Settings, and IaC tools.

use provisioning_mcp_server::tools::settings::{DeploymentMode, SettingsTools};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::Mutex;

/// Tool execution result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolExecution {
    pub tool_name: String,
    pub result: Value,
    pub duration_ms: u64,
}

/// MCP tool registry for provisioning system
pub struct ToolRegistry {
    tools: std::collections::HashMap<String, ToolDefinition>,
    settings_tools: Mutex<SettingsTools>,
}

/// Tool definition for MCP
#[derive(Debug, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub category: ToolCategory,
    pub input_schema: Value,
}

/// Tool categories
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolCategory {
    Rag,
    Guidance,
    Settings,
    Iac,
}

impl ToolRegistry {
    /// Create a new tool registry
    pub fn new() -> Self {
        let mut registry = Self {
            tools: std::collections::HashMap::new(),
            settings_tools: Mutex::new(SettingsTools::new()),
        };
        registry.register_all_tools();
        registry
    }

    /// Register all tool categories (RAG, Guidance, Settings, IaC)
    fn register_all_tools(&mut self) {
        self.register_rag_tools();
        self.register_guidance_tools();
        self.register_settings_tools();
        self.register_iac_tools();
    }

    /// Register RAG tools
    fn register_rag_tools(&mut self) {
        self.tools.insert(
            "rag_ask_question".to_string(),
            ToolDefinition {
                name: "rag_ask_question".to_string(),
                description: "Ask a question using RAG (Retrieval-Augmented Generation) with knowledge base search".to_string(),
                category: ToolCategory::Rag,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "question": {"type": "string", "description": "The question to ask"},
                        "context": {"type": "string", "description": "Optional context for the question"},
                        "top_k": {"type": "integer", "description": "Number of top results to return", "default": 3}
                    },
                    "required": ["question"]
                }),
            },
        );

        self.tools.insert(
            "rag_semantic_search".to_string(),
            ToolDefinition {
                name: "rag_semantic_search".to_string(),
                description: "Perform semantic search on the knowledge base".to_string(),
                category: ToolCategory::Rag,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Search query"},
                        "category": {"type": "string", "description": "Optional category filter"},
                        "top_k": {"type": "integer", "description": "Number of results", "default": 5}
                    },
                    "required": ["query"]
                }),
            },
        );

        self.tools.insert(
            "rag_get_status".to_string(),
            ToolDefinition {
                name: "rag_get_status".to_string(),
                description: "Get the current status of the RAG knowledge base".to_string(),
                category: ToolCategory::Rag,
                input_schema: json!({
                    "type": "object",
                    "properties": {}
                }),
            },
        );
    }

    /// Register Guidance tools
    fn register_guidance_tools(&mut self) {
        self.tools.insert(
            "guidance_check_system_status".to_string(),
            ToolDefinition {
                name: "guidance_check_system_status".to_string(),
                description: "Check the current system status and configuration".to_string(),
                category: ToolCategory::Guidance,
                input_schema: json!({
                    "type": "object",
                    "properties": {}
                }),
            },
        );

        self.tools.insert(
            "guidance_suggest_next_action".to_string(),
            ToolDefinition {
                name: "guidance_suggest_next_action".to_string(),
                description: "Get suggestions for the next action based on current system state".to_string(),
                category: ToolCategory::Guidance,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "context": {"type": "string", "description": "Optional context for suggestion"}
                    }
                }),
            },
        );

        self.tools.insert(
            "guidance_find_docs".to_string(),
            ToolDefinition {
                name: "guidance_find_docs".to_string(),
                description: "Find relevant documentation and guides".to_string(),
                category: ToolCategory::Guidance,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "What to search for"},
                        "context": {"type": "string", "description": "Optional context"}
                    },
                    "required": ["query"]
                }),
            },
        );

        self.tools.insert(
            "guidance_troubleshoot".to_string(),
            ToolDefinition {
                name: "guidance_troubleshoot".to_string(),
                description: "Troubleshoot an issue or error".to_string(),
                category: ToolCategory::Guidance,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "error": {"type": "string", "description": "Error message or description"},
                        "context": {"type": "string", "description": "Context about the issue"}
                    },
                    "required": ["error"]
                }),
            },
        );

        self.tools.insert(
            "guidance_validate_config".to_string(),
            ToolDefinition {
                name: "guidance_validate_config".to_string(),
                description: "Validate a configuration file or settings".to_string(),
                category: ToolCategory::Guidance,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "config_path": {"type": "string", "description": "Path to configuration file"}
                    },
                    "required": ["config_path"]
                }),
            },
        );
    }

    /// Register Settings tools
    fn register_settings_tools(&mut self) {
        self.tools.insert(
            "installer_get_settings".to_string(),
            ToolDefinition {
                name: "installer_get_settings".to_string(),
                description: "Get installer settings and configuration".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Optional settings query"}
                    }
                }),
            },
        );

        self.tools.insert(
            "installer_complete_config".to_string(),
            ToolDefinition {
                name: "installer_complete_config".to_string(),
                description: "Complete partial configuration with defaults".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "config": {"type": "object", "description": "Partial configuration"}
                    }
                }),
            },
        );

        self.tools.insert(
            "installer_validate_config".to_string(),
            ToolDefinition {
                name: "installer_validate_config".to_string(),
                description: "Validate configuration against schema".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "config": {"type": "object", "description": "Configuration to validate"}
                    },
                    "required": ["config"]
                }),
            },
        );

        self.tools.insert(
            "installer_get_defaults".to_string(),
            ToolDefinition {
                name: "installer_get_defaults".to_string(),
                description: "Get default settings for a deployment mode".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "mode": {"type": "string", "description": "Deployment mode"}
                    },
                    "required": ["mode"]
                }),
            },
        );

        self.tools.insert(
            "installer_platform_recommendations".to_string(),
            ToolDefinition {
                name: "installer_platform_recommendations".to_string(),
                description: "Get platform-specific recommendations".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {}
                }),
            },
        );

        self.tools.insert(
            "installer_service_recommendations".to_string(),
            ToolDefinition {
                name: "installer_service_recommendations".to_string(),
                description: "Get service recommendations for a deployment mode".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "mode": {"type": "string", "description": "Deployment mode"}
                    },
                    "required": ["mode"]
                }),
            },
        );

        self.tools.insert(
            "installer_resource_recommendations".to_string(),
            ToolDefinition {
                name: "installer_resource_recommendations".to_string(),
                description: "Get resource recommendations for a deployment mode".to_string(),
                category: ToolCategory::Settings,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "mode": {"type": "string", "description": "Deployment mode"}
                    },
                    "required": ["mode"]
                }),
            },
        );
    }

    /// Register IaC tools
    fn register_iac_tools(&mut self) {
        self.tools.insert(
            "iac_detect_technologies".to_string(),
            ToolDefinition {
                name: "iac_detect_technologies".to_string(),
                description: "Detect technologies used in infrastructure".to_string(),
                category: ToolCategory::Iac,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Path to analyze"}
                    }
                }),
            },
        );

        self.tools.insert(
            "iac_analyze_completeness".to_string(),
            ToolDefinition {
                name: "iac_analyze_completeness".to_string(),
                description: "Analyze completeness of infrastructure configuration".to_string(),
                category: ToolCategory::Iac,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "config": {"type": "object", "description": "Configuration to analyze"}
                    }
                }),
            },
        );

        self.tools.insert(
            "iac_infer_requirements".to_string(),
            ToolDefinition {
                name: "iac_infer_requirements".to_string(),
                description: "Infer infrastructure requirements from description".to_string(),
                category: ToolCategory::Iac,
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "description": {"type": "string", "description": "Infrastructure description"}
                    },
                    "required": ["description"]
                }),
            },
        );
    }

    /// Get all tool definitions
    pub fn list_tools(&self) -> Vec<ToolDefinition> {
        self.tools.values().cloned().collect()
    }

    /// Get tools by category
    pub fn tools_by_category(&self, category: ToolCategory) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .filter(|t| t.category == category)
            .cloned()
            .collect()
    }

    /// Check if tool exists
    pub fn has_tool(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }

    /// Execute a tool (async)
    pub async fn execute(&self, tool_name: &str, args: &Value) -> Result<Value, String> {
        match tool_name {
            // RAG tools
            "rag_ask_question" => self.rag_ask_question(args).await,
            "rag_semantic_search" => self.rag_semantic_search(args).await,
            "rag_get_status" => self.rag_get_status(args).await,
            // Guidance tools
            "guidance_check_system_status" => self.guidance_check_system_status(args).await,
            "guidance_suggest_next_action" => self.guidance_suggest_next_action(args).await,
            "guidance_find_docs" => self.guidance_find_docs(args).await,
            "guidance_troubleshoot" => self.guidance_troubleshoot(args).await,
            "guidance_validate_config" => self.guidance_validate_config(args).await,
            // Settings tools
            "installer_get_settings" => self.installer_get_settings(args).await,
            "installer_complete_config" => self.installer_complete_config(args).await,
            "installer_validate_config" => self.installer_validate_config(args).await,
            "installer_get_defaults" => self.installer_get_defaults(args).await,
            "installer_platform_recommendations" => {
                self.installer_platform_recommendations(args).await
            }
            "installer_service_recommendations" => {
                self.installer_service_recommendations(args).await
            }
            "installer_resource_recommendations" => {
                self.installer_resource_recommendations(args).await
            }
            // IaC tools
            "iac_detect_technologies" => self.iac_detect_technologies(args).await,
            "iac_analyze_completeness" => self.iac_analyze_completeness(args).await,
            "iac_infer_requirements" => self.iac_infer_requirements(args).await,
            _ => Err(format!("Unknown tool: {}", tool_name)),
        }
    }

    // ========== RAG TOOL IMPLEMENTATIONS ==========

    async fn rag_ask_question(&self, args: &Value) -> Result<Value, String> {
        let question = args
            .get("question")
            .and_then(|v| v.as_str())
            .ok_or("question parameter required")?;
        Ok(json!({
            "status": "success",
            "tool": "rag_ask_question",
            "message": format!("RAG query would be processed for: {}", question)
        }))
    }

    async fn rag_semantic_search(&self, args: &Value) -> Result<Value, String> {
        let query = args
            .get("query")
            .and_then(|v| v.as_str())
            .ok_or("query parameter required")?;
        Ok(json!({
            "status": "success",
            "tool": "rag_semantic_search",
            "message": format!("Semantic search would be performed for: {}", query),
            "results": []
        }))
    }

    async fn rag_get_status(&self, _args: &Value) -> Result<Value, String> {
        Ok(json!({
            "status": "active",
            "tool": "rag_get_status",
            "knowledge_base": {
                "documents_loaded": true,
                "total_documents": 76,
                "categories": ["architecture", "deployment", "security", "reliability"]
            }
        }))
    }

    // ========== GUIDANCE TOOL IMPLEMENTATIONS ==========

    /// Execute a Nushell command and parse JSON output
    async fn execute_nu_command(cmd: &str) -> Result<Value, String> {
        use tokio::process::Command;

        let output = Command::new("nu")
            .arg("-c")
            .arg(cmd)
            .output()
            .await
            .map_err(|e| format!("Failed to execute Nushell: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Nushell command failed: {}", stderr));
        }

        serde_json::from_slice(&output.stdout)
            .map_err(|e| format!("Failed to parse JSON output: {}", e))
    }

    async fn guidance_check_system_status(&self, _args: &Value) -> Result<Value, String> {
        Self::execute_nu_command("provisioning status-json").await
    }

    async fn guidance_suggest_next_action(&self, _args: &Value) -> Result<Value, String> {
        Self::execute_nu_command("provisioning next").await
    }

    async fn guidance_find_docs(&self, args: &Value) -> Result<Value, String> {
        let query = args
            .get("query")
            .and_then(|v| v.as_str())
            .ok_or("query parameter required")?;

        let cmd = format!("provisioning guide {}", query);
        Self::execute_nu_command(&cmd).await
    }

    async fn guidance_troubleshoot(&self, _args: &Value) -> Result<Value, String> {
        Self::execute_nu_command("provisioning health-json").await
    }

    async fn guidance_validate_config(&self, args: &Value) -> Result<Value, String> {
        let config_path = args
            .get("config_path")
            .and_then(|v| v.as_str())
            .ok_or("config_path parameter required")?;

        let cmd = format!("validate-system-config {}", config_path);
        Self::execute_nu_command(&cmd).await
    }

    // ========== SETTINGS TOOL IMPLEMENTATIONS ==========

    async fn installer_get_settings(&self, args: &Value) -> Result<Value, String> {
        let query = args.get("query").and_then(|v| v.as_str());

        let mut settings_tools = self.settings_tools.lock().await;

        let settings = settings_tools
            .get_settings(query)
            .await
            .map_err(|e| format!("Failed to get settings: {}", e))?;

        Ok(json!({
            "status": "success",
            "tool": "installer_get_settings",
            "platforms": settings.get("platforms"),
            "modes": settings.get("modes"),
            "default_domain": settings.get("default_domain"),
            "auto_generate_secrets": settings.get("auto_generate_secrets"),
            "available_services": settings.get("available_services")
        }))
    }

    async fn installer_complete_config(&self, args: &Value) -> Result<Value, String> {
        let config = args
            .get("config")
            .cloned()
            .ok_or("config parameter required")?;

        let mut settings_tools = self.settings_tools.lock().await;

        settings_tools
            .complete_config(config)
            .await
            .map_err(|e| format!("Failed to complete config: {}", e))
    }

    async fn installer_validate_config(&self, args: &Value) -> Result<Value, String> {
        let config = args
            .get("config")
            .cloned()
            .ok_or("config parameter required")?;

        let settings_tools = self.settings_tools.lock().await;

        settings_tools
            .validate_config(config)
            .map_err(|e| format!("Failed to validate config: {}", e))
    }

    async fn installer_get_defaults(&self, args: &Value) -> Result<Value, String> {
        let mode = args
            .get("mode")
            .and_then(|v| v.as_str())
            .ok_or("mode parameter required")?;

        let settings_tools = self.settings_tools.lock().await;

        let defaults = settings_tools
            .get_mode_defaults(mode)
            .map_err(|e| format!("Failed to get defaults: {}", e))?;

        Ok(json!({
            "status": "success",
            "tool": "installer_get_defaults",
            "mode": defaults.get("mode"),
            "description": defaults.get("description"),
            "service_count": defaults.get("service_count"),
            "min_cpu_cores": defaults.get("min_cpu_cores"),
            "min_memory_gb": defaults.get("min_memory_gb"),
            "recommended_services": defaults.get("recommended_services"),
            "auto_generate_secrets": defaults.get("auto_generate_secrets"),
            "domain": defaults.get("domain")
        }))
    }

    async fn installer_platform_recommendations(&self, _args: &Value) -> Result<Value, String> {
        let mut settings_tools = self.settings_tools.lock().await;

        let recommendations = settings_tools
            .get_platform_recommendations()
            .await
            .map_err(|e| format!("Failed to get platform recommendations: {}", e))?;

        Ok(json!({
            "status": "success",
            "tool": "installer_platform_recommendations",
            "recommendations": recommendations
        }))
    }

    async fn installer_service_recommendations(&self, args: &Value) -> Result<Value, String> {
        let mode_str = args
            .get("mode")
            .and_then(|v| v.as_str())
            .ok_or("mode parameter required")?;

        let mode = DeploymentMode::from_str(mode_str)
            .ok_or(format!("Invalid deployment mode: {}", mode_str))?;

        let settings_tools = self.settings_tools.lock().await;

        let recommendations = settings_tools.get_service_recommendations(&mode);

        Ok(json!({
            "status": "success",
            "tool": "installer_service_recommendations",
            "mode": mode_str,
            "recommendations": recommendations
        }))
    }

    async fn installer_resource_recommendations(&self, args: &Value) -> Result<Value, String> {
        let mode_str = args
            .get("mode")
            .and_then(|v| v.as_str())
            .ok_or("mode parameter required")?;

        let mode = DeploymentMode::from_str(mode_str)
            .ok_or(format!("Invalid deployment mode: {}", mode_str))?;

        let settings_tools = self.settings_tools.lock().await;

        let recommendations = settings_tools.get_resource_recommendations(&mode);

        Ok(json!({
            "status": "success",
            "tool": "installer_resource_recommendations",
            "mode": mode_str,
            "recommendations": recommendations
        }))
    }

    // ========== IAC TOOL IMPLEMENTATIONS ==========

    async fn iac_detect_technologies(&self, args: &Value) -> Result<Value, String> {
        let path = args
            .get("path")
            .and_then(|v| v.as_str())
            .ok_or("path parameter required")?;

        let path_obj = std::path::Path::new(path);
        let mut technologies = Vec::new();

        // Check for Kubernetes
        if path_obj.join("kustomization.yaml").exists() || path_obj.join("deployment.yaml").exists()
        {
            technologies.push("kubernetes");
        }

        // Check for Docker
        if path_obj.join("Dockerfile").exists() || path_obj.join("docker-compose.yml").exists() {
            technologies.push("docker");
        }

        // Check for Terraform
        if path_obj.join("main.tf").exists() {
            technologies.push("terraform");
        }

        // Check for KCL (legacy)
        if path_obj.join("kcl.mod").exists() {
            technologies.push("kcl");
        }

        // Check for Nickel (current IaC)
        if path_obj.join("main.ncl").exists() {
            technologies.push("nickel");
        }

        Ok(json!({
            "status": "success",
            "tool": "iac_detect_technologies",
            "path": path,
            "technologies": technologies
        }))
    }

    async fn iac_analyze_completeness(&self, args: &Value) -> Result<Value, String> {
        let path = args
            .get("path")
            .and_then(|v| v.as_str())
            .ok_or("path parameter required")?;

        let path_obj = std::path::Path::new(path);
        let mut missing = Vec::new();

        // Check for essential infrastructure files
        if !path_obj.join("infrastructure.ncl").exists() {
            missing.push("infrastructure.ncl");
        }
        if !path_obj.join("config.toml").exists() {
            missing.push("config.toml");
        }
        if !path_obj.join("README.md").exists() {
            missing.push("README.md");
        }

        let complete = missing.is_empty();
        let completeness_score = if missing.is_empty() { 1.0 } else { 0.7 };

        Ok(json!({
            "status": "success",
            "tool": "iac_analyze_completeness",
            "complete": complete,
            "completeness_score": completeness_score,
            "missing_files": missing
        }))
    }

    async fn iac_infer_requirements(&self, args: &Value) -> Result<Value, String> {
        let _description = args
            .get("description")
            .and_then(|v| v.as_str())
            .ok_or("description parameter required")?;

        // Basic requirements inference (can be enhanced with ML later)
        Ok(json!({
            "status": "success",
            "tool": "iac_infer_requirements",
            "requirements": {
                "compute": "2 CPU, 4GB RAM (minimum)",
                "storage": "20GB",
                "network": "Private network recommended",
                "high_availability": false
            }
        }))
    }
}

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