//! RAG Agent with tool calling integration
//!
//! Extends RagAgent with Claude function calling support, allowing the agent
//! to:
//! - Understand and invoke available tools
//! - Handle tool results and integrate them into responses
//! - Manage tool execution with security and rate limiting
//! - Track tool usage statistics

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use tracing::info;

use crate::agent::{AgentResponse, RagAgent};
use crate::error::Result;
use crate::tools::{RagTool, ToolInput, ToolOutput, ToolRegistry};

/// Agent response with tool usage information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResponseWithTools {
    /// The answer to the user's question
    pub answer: String,

    /// Tools that were invoked
    pub tools_used: Vec<String>,

    /// Tool execution results
    pub tool_results: Vec<(String, ToolOutput)>,

    /// Original agent response
    pub base_response: AgentResponse,
}

/// Tool call request from Claude
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Tool name to invoke
    pub name: String,

    /// Tool input parameters
    pub input: ToolInput,
}

/// Tool-enabled RAG agent
pub struct ToolEnabledRagAgent {
    /// Base RAG agent
    agent: RagAgent,

    /// Tool registry
    tool_registry: Arc<ToolRegistry>,

    /// User ID for access control
    user_id: String,
}

impl ToolEnabledRagAgent {
    /// Create a new tool-enabled RAG agent
    pub fn new(agent: RagAgent, tool_registry: Arc<ToolRegistry>, user_id: String) -> Self {
        Self {
            agent,
            tool_registry,
            user_id,
        }
    }

    /// Ask a question with potential tool invocation
    pub async fn ask_with_tools(&self, question: &str) -> Result<AgentResponseWithTools> {
        // 1. Get base response from RAG agent
        let base_response = self.agent.ask(question).await?;

        // 2. Analyze response for tool calls (in production, Claude would suggest
        //    these)
        let tool_calls = self.analyze_for_tool_calls(&base_response.answer).await;

        info!("Identified {} potential tool calls", tool_calls.len());

        // 3. Execute tool calls with security validation
        let mut tool_results = Vec::new();
        let mut tools_used = Vec::new();

        for call in tool_calls {
            match self
                .tool_registry
                .call_tool(&call.name, call.input, &self.user_id)
                .await
            {
                Ok(output) => {
                    info!("Tool '{}' executed successfully", call.name);
                    tool_results.push((call.name.clone(), output));
                    tools_used.push(call.name);
                }
                Err(e) => {
                    info!("Tool '{}' failed: {}", call.name, e);
                    let error_output = ToolOutput {
                        result: format!("Tool execution failed: {}", e),
                        success: false,
                        error: Some(e.to_string()),
                    };
                    tool_results.push((call.name.clone(), error_output));
                }
            }
        }

        // 4. Integrate tool results into response
        let enhanced_answer = if !tool_results.is_empty() {
            self.integrate_tool_results(&base_response.answer, &tool_results)
        } else {
            base_response.answer.clone()
        };

        Ok(AgentResponseWithTools {
            answer: enhanced_answer,
            tools_used,
            tool_results,
            base_response,
        })
    }

    /// Analyze response text for tool invocations
    /// In production, Claude would suggest tool calls directly
    async fn analyze_for_tool_calls(&self, _answer: &str) -> Vec<ToolCall> {
        // For now, return empty - in production, this would:
        // 1. Parse Claude's response for tool_use tags
        // 2. Extract tool names and parameters
        // 3. Return structured ToolCall objects

        // Example of what production would look like:
        // If answer contains "<tool_use name="create_server"...>
        // Extract parameters and return as ToolCall

        Vec::new()
    }

    /// Integrate tool results into the answer
    fn integrate_tool_results(
        &self,
        answer: &str,
        tool_results: &[(String, ToolOutput)],
    ) -> String {
        let mut enhanced = format!("{}\n\n## Tool Execution Results:\n", answer);

        for (tool_name, output) in tool_results {
            enhanced.push_str(&format!(
                "\n**{}**: {}\n",
                tool_name,
                if output.success { "✓" } else { "✗" }
            ));
            enhanced.push_str(&format!("Result: {}\n", output.result));

            if let Some(error) = &output.error {
                enhanced.push_str(&format!("Error: {}\n", error));
            }
        }

        enhanced
    }

    /// Get available tools as definitions for Claude
    pub async fn get_available_tools(&self) -> Vec<String> {
        self.tool_registry
            .get_definitions()
            .await
            .into_iter()
            .map(|def| def.name)
            .collect()
    }

    /// Get tool statistics
    pub async fn get_tool_stats(&self) -> crate::tools::ToolCallStats {
        self.tool_registry.get_stats().await
    }

    /// Register a new tool
    pub async fn register_tool(&self, tool: Arc<dyn RagTool>) -> Result<()> {
        self.tool_registry.register(tool).await
    }

    /// Get the underlying RAG agent
    pub fn base_agent(&self) -> &RagAgent {
        &self.agent
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{CreateServerTool, WorkspaceStatusTool};

    async fn create_test_registry() -> Arc<ToolRegistry> {
        // Note: In real tests, we'd create a full RagAgent with all dependencies
        // For unit tests, we'll focus on tool registry functionality
        let registry = Arc::new(ToolRegistry::new());
        registry
            .register(Arc::new(CreateServerTool::new()))
            .await
            .unwrap();
        registry
            .register(Arc::new(WorkspaceStatusTool::new(
                "test-workspace".to_string(),
            )))
            .await
            .unwrap();

        registry
    }

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

        let tools = registry.get_definitions().await;
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0].name, "get_workspace_status");
    }

    #[tokio::test]
    async fn test_tool_output_integration() {
        let results = vec![(
            "test_tool".to_string(),
            ToolOutput {
                result: "Success".to_string(),
                success: true,
                error: None,
            },
        )];

        let answer = "Initial answer";
        let mut enhanced = format!("{}\n\n## Tool Execution Results:\n", answer);

        for (tool_name, output) in &results {
            enhanced.push_str(&format!(
                "\n**{}**: {}\n",
                tool_name,
                if output.success { "✓" } else { "✗" }
            ));
            enhanced.push_str(&format!("Result: {}\n", output.result));
        }

        assert!(enhanced.contains("Initial answer"));
        assert!(enhanced.contains("test_tool"));
        assert!(enhanced.contains("Success"));
    }

    #[test]
    fn test_tool_call_structure() {
        let call = ToolCall {
            name: "create_server".to_string(),
            input: ToolInput {
                params: serde_json::json!({
                    "hostname": "web-01",
                    "cores": 4
                }),
            },
        };

        assert_eq!(call.name, "create_server");
        assert_eq!(call.input.params["hostname"], "web-01");
    }

    #[test]
    fn test_agent_response_with_tools() {
        use crate::AgentResponse;

        let response = AgentResponseWithTools {
            answer: "Test answer".to_string(),
            tools_used: vec!["tool1".to_string()],
            tool_results: vec![(
                "tool1".to_string(),
                ToolOutput {
                    result: "Result".to_string(),
                    success: true,
                    error: None,
                },
            )],
            base_response: AgentResponse {
                answer: "Base answer".to_string(),
                sources: vec![],
                confidence: 0.8,
                context: "Test context".to_string(),
            },
        };

        assert_eq!(response.tools_used.len(), 1);
        assert_eq!(response.tool_results.len(), 1);
        assert_eq!(response.base_response.confidence, 0.8);
    }
}
