dynamic_grounding_for_github_copilot 0.1.0

//! MCP tool handlers that expose Gemini API operations via the Model Context Protocol.
//!
//! This module implements four tools:
//! - search_codebase: Natural language search across code files
//! - analyze_files: Cross-file analysis and relationship discovery
//! - ask_about_code: Context-based Q&A about code
//! - summarize_directory: High-level overview of project structure

use crate::gemini::{FileContent, GeminiClient};
use anyhow::Result;
use mocopr_core::ToolExecutor;
use mocopr_core::types::{Content, TextContent, ToolsCallResponse};
use mocopr_macros::Tool;
use serde_json::{Value, json};
use std::sync::Arc;

/// Tool for searching codebase using natural language queries
#[derive(Tool, Clone)]
#[tool(
    name = "search_codebase",
    description = "Search through code files using natural language queries. Analyzes code using Google Gemini AI to find relevant sections, patterns, and implementations. Parameters: 'query' (string, required) - the search query in natural language; 'files' (array, required) - array of objects with 'path' and 'content' fields containing the code to search through."
)]
pub struct SearchCodebaseTool {
    gemini_client: Arc<GeminiClient>,
}

impl SearchCodebaseTool {
    pub fn new(gemini_client: Arc<GeminiClient>) -> Self {
        Self { gemini_client }
    }

    async fn execute_impl(&self, args: Value) -> Result<Value> {
        let query = args
            .get("query")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: query"))?;

        let files = args
            .get("files")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: files (array)"))?;

        // Parse file objects
        let mut file_contents = Vec::new();
        for file_obj in files {
            let path = file_obj
                .get("path")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'path' field"))?;

            let content = file_obj
                .get("content")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'content' field"))?;

            file_contents.push(FileContent {
                path: path.to_string(),
                content: content.to_string(),
            });
        }

        if file_contents.is_empty() {
            return Err(anyhow::anyhow!("files array cannot be empty"));
        }

        // Call Gemini API
        let response = self
            .gemini_client
            .search_codebase(query, &file_contents)
            .await
            .map_err(|e| anyhow::anyhow!("Gemini API error: {}", e))?;

        Ok(json!({
            "query": query,
            "files_searched": file_contents.len(),
            "results": response.text,
            "quota_status": response.quota_status.format_message()
        }))
    }
}

#[async_trait::async_trait]
impl ToolExecutor for SearchCodebaseTool {
    async fn execute(&self, arguments: Option<Value>) -> mocopr_core::Result<ToolsCallResponse> {
        let args = arguments.unwrap_or_default();

        match self.execute_impl(args).await {
            Ok(result) => Ok(ToolsCallResponse::success(vec![Content::Text(
                TextContent::new(result.to_string()),
            )])),
            Err(e) => Ok(ToolsCallResponse::error(vec![Content::Text(
                TextContent::new(e.to_string()),
            )])),
        }
    }
}

/// Tool for analyzing multiple files and discovering relationships
#[derive(Tool, Clone)]
#[tool(
    name = "analyze_files",
    description = "Analyze multiple code files to understand relationships, dependencies, and interactions between components using Google Gemini AI. Parameters: 'question' (string, required) - what you want to learn about the files; 'files' (array, required) - array of objects with 'path' and 'content' fields containing the code to analyze."
)]
pub struct AnalyzeFilesTool {
    gemini_client: Arc<GeminiClient>,
}

impl AnalyzeFilesTool {
    pub fn new(gemini_client: Arc<GeminiClient>) -> Self {
        Self { gemini_client }
    }

    async fn execute_impl(&self, args: Value) -> Result<Value> {
        let question = args
            .get("question")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: question"))?;

        let files = args
            .get("files")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: files (array)"))?;

        // Parse file objects
        let mut file_contents = Vec::new();
        for file_obj in files {
            let path = file_obj
                .get("path")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'path' field"))?;

            let content = file_obj
                .get("content")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'content' field"))?;

            file_contents.push(FileContent {
                path: path.to_string(),
                content: content.to_string(),
            });
        }

        if file_contents.is_empty() {
            return Err(anyhow::anyhow!("files array cannot be empty"));
        }

        // Call Gemini API
        let response = self
            .gemini_client
            .analyze_files(&file_contents, question)
            .await
            .map_err(|e| anyhow::anyhow!("Gemini API error: {}", e))?;

        Ok(json!({
            "question": question,
            "files_analyzed": file_contents.len(),
            "analysis": response.text,
            "quota_status": response.quota_status.format_message()
        }))
    }
}

#[async_trait::async_trait]
impl ToolExecutor for AnalyzeFilesTool {
    async fn execute(&self, arguments: Option<Value>) -> mocopr_core::Result<ToolsCallResponse> {
        let args = arguments.unwrap_or_default();

        match self.execute_impl(args).await {
            Ok(result) => Ok(ToolsCallResponse::success(vec![Content::Text(
                TextContent::new(result.to_string()),
            )])),
            Err(e) => Ok(ToolsCallResponse::error(vec![Content::Text(
                TextContent::new(e.to_string()),
            )])),
        }
    }
}

/// Tool for asking questions about specific code with context
#[derive(Tool, Clone)]
#[tool(
    name = "ask_about_code",
    description = "Ask questions about specific code sections with Google Gemini AI providing detailed explanations, design patterns, and implementation details. Parameters: 'context' (string, required) - the code or description to ask about; 'question' (string, required) - your specific question about the code."
)]
pub struct AskAboutCodeTool {
    gemini_client: Arc<GeminiClient>,
}

impl AskAboutCodeTool {
    pub fn new(gemini_client: Arc<GeminiClient>) -> Self {
        Self { gemini_client }
    }

    async fn execute_impl(&self, args: Value) -> Result<Value> {
        let context = args
            .get("context")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: context"))?;

        let question = args
            .get("question")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: question"))?;

        // Call Gemini API
        let response = self
            .gemini_client
            .ask_about_code(context, question)
            .await
            .map_err(|e| anyhow::anyhow!("Gemini API error: {}", e))?;

        Ok(json!({
            "question": question,
            "answer": response.text,
            "quota_status": response.quota_status.format_message()
        }))
    }
}

#[async_trait::async_trait]
impl ToolExecutor for AskAboutCodeTool {
    async fn execute(&self, arguments: Option<Value>) -> mocopr_core::Result<ToolsCallResponse> {
        let args = arguments.unwrap_or_default();

        match self.execute_impl(args).await {
            Ok(result) => Ok(ToolsCallResponse::success(vec![Content::Text(
                TextContent::new(result.to_string()),
            )])),
            Err(e) => Ok(ToolsCallResponse::error(vec![Content::Text(
                TextContent::new(e.to_string()),
            )])),
        }
    }
}

/// Tool for summarizing directory structure and contents
#[derive(Tool, Clone)]
#[tool(
    name = "summarize_directory",
    description = "Get a high-level overview of a directory's structure and key files using Google Gemini AI to understand project organization and architecture. Parameters: 'directory_structure' (string, required) - text representation of the directory tree; 'files' (array, required) - array of objects with 'path' and 'content' fields for key files to summarize."
)]
pub struct SummarizeDirectoryTool {
    gemini_client: Arc<GeminiClient>,
}

impl SummarizeDirectoryTool {
    pub fn new(gemini_client: Arc<GeminiClient>) -> Self {
        Self { gemini_client }
    }

    async fn execute_impl(&self, args: Value) -> Result<Value> {
        let directory_structure = args
            .get("directory_structure")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: directory_structure"))?;

        let files = args
            .get("files")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: files (array)"))?;

        // Parse file objects
        let mut file_contents = Vec::new();
        for file_obj in files {
            let path = file_obj
                .get("path")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'path' field"))?;

            let content = file_obj
                .get("content")
                .and_then(|v| v.as_str())
                .ok_or_else(|| anyhow::anyhow!("Each file must have a 'content' field"))?;

            file_contents.push(FileContent {
                path: path.to_string(),
                content: content.to_string(),
            });
        }

        // Call Gemini API
        let response = self
            .gemini_client
            .summarize_directory(directory_structure, &file_contents)
            .await
            .map_err(|e| anyhow::anyhow!("Gemini API error: {}", e))?;

        Ok(json!({
            "directory": directory_structure,
            "files_included": file_contents.len(),
            "summary": response.text,
            "quota_status": response.quota_status.format_message()
        }))
    }
}

#[async_trait::async_trait]
impl ToolExecutor for SummarizeDirectoryTool {
    async fn execute(&self, arguments: Option<Value>) -> mocopr_core::Result<ToolsCallResponse> {
        let args = arguments.unwrap_or_default();

        match self.execute_impl(args).await {
            Ok(result) => Ok(ToolsCallResponse::success(vec![Content::Text(
                TextContent::new(result.to_string()),
            )])),
            Err(e) => Ok(ToolsCallResponse::error(vec![Content::Text(
                TextContent::new(e.to_string()),
            )])),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::api_key::ApiKeyProvider;
    use crate::quota::QuotaTracker;

    struct MockApiKeyProvider;

    #[async_trait::async_trait]
    impl ApiKeyProvider for MockApiKeyProvider {
        async fn get_key(&self) -> crate::error::Result<crate::api_key::SecureString> {
            Ok(crate::api_key::SecureString::new("test_key".to_string()))
        }
    }

    #[test]
    fn test_tool_creation() {
        let provider = Arc::new(MockApiKeyProvider) as Arc<dyn ApiKeyProvider>;
        let quota_tracker = Arc::new(QuotaTracker::new());
        let client = Arc::new(GeminiClient::new(provider, quota_tracker));

        let _search_tool = SearchCodebaseTool::new(client.clone());
        let _analyze_tool = AnalyzeFilesTool::new(client.clone());
        let _ask_tool = AskAboutCodeTool::new(client.clone());
        let _summarize_tool = SummarizeDirectoryTool::new(client);
    }
}