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

//! Core AI service implementation with RAG integration

use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, info};

/// MCP tool invocation request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpToolRequest {
    /// Tool name (e.g., "execute_provisioning_plan")
    pub tool_name: String,
    /// Tool arguments as JSON
    pub args: serde_json::Value,
}

/// MCP tool invocation response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpToolResponse {
    /// Tool execution result
    pub result: serde_json::Value,
    /// Execution time in milliseconds
    pub duration_ms: u64,
}

/// RAG-powered question request with optional hybrid tool execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AskRequest {
    /// User question
    pub question: String,
    /// Optional context
    pub context: Option<String>,
    /// Enable automatic tool execution (hybrid mode)
    /// When true, the RAG answer may trigger tool calls automatically
    /// Default: false (explicit tool calls only)
    #[serde(default)]
    pub enable_tool_execution: Option<bool>,
    /// Maximum number of tools to execute automatically
    /// Default: 3
    #[serde(default)]
    pub max_tool_calls: Option<u32>,
}

/// RAG-powered question response with optional tool execution results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AskResponse {
    /// Answer from AI
    pub answer: String,
    /// Source documents used
    pub sources: Vec<String>,
    /// Confidence level (0-100)
    pub confidence: u8,
    /// Reasoning explanation
    pub reasoning: String,
    /// Tool executions performed (if hybrid mode enabled)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_executions: Option<Vec<crate::mcp::ToolExecution>>,
}

/// Extension DAG node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DagNode {
    /// Extension/component name
    pub name: String,
    /// Dependencies on other nodes
    pub dependencies: Vec<String>,
    /// Component version
    pub version: String,
}

/// Extension DAG response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DagResponse {
    /// DAG nodes (extensions)
    pub nodes: Vec<DagNode>,
    /// DAG edges (dependencies)
    pub edges: Vec<(String, String)>,
}

/// Best practice entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BestPractice {
    /// Practice title
    pub title: String,
    /// Practice description
    pub description: String,
    /// Category (e.g., "deployment", "security")
    pub category: String,
    /// Relevance score (0-100)
    pub relevance: u8,
}

/// Knowledge base document (from RAG ingestion)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeDocument {
    /// Document ID
    pub id: String,
    /// Document type
    #[serde(rename = "type")]
    pub doc_type: String,
    /// Document title
    pub title: Option<String>,
    /// Document name (for extensions)
    pub name: Option<String>,
    /// Full content
    pub content: String,
    /// Document category
    pub category: String,
    /// Tags
    pub tags: Vec<String>,
    /// Relevance/importance
    pub relevance: Option<u8>,
    /// Dependencies (for extensions)
    pub dependencies: Option<Vec<String>>,
}

/// Knowledge base with documents and relationships
#[derive(Debug, Clone)]
pub struct KnowledgeBase {
    /// All documents indexed by ID
    pub documents: std::collections::HashMap<String, KnowledgeDocument>,
    /// Document relationships
    pub relationships: Vec<Relationship>,
}

/// Document relationship
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Relationship {
    /// Source document ID
    pub source_id: String,
    /// Target document ID
    pub target_id: String,
    /// Relationship type
    pub relationship_type: String,
    /// Strength (0-1)
    pub strength: f32,
}

/// Core AI service
pub struct AiService {
    /// Server address
    addr: SocketAddr,
    /// Knowledge base
    knowledge_base: Arc<RwLock<KnowledgeBase>>,
    /// MCP tool registry
    tool_registry: crate::mcp::ToolRegistry,
}

impl AiService {
    /// Create new AI service
    pub fn new(addr: SocketAddr) -> Self {
        Self {
            addr,
            knowledge_base: Arc::new(RwLock::new(KnowledgeBase {
                documents: std::collections::HashMap::new(),
                relationships: Vec::new(),
            })),
            tool_registry: crate::mcp::ToolRegistry::new(),
        }
    }

    /// Load knowledge base from JSON files
    pub async fn load_knowledge_base(&self, knowledge_base_dir: &str) -> Result<()> {
        info!("Loading knowledge base from: {}", knowledge_base_dir);

        // Load best practice documents
        let bp_path = format!("{}/best-practices-docs.json", knowledge_base_dir);
        let bp_content = std::fs::read_to_string(&bp_path)
            .map_err(|e| anyhow::anyhow!("Failed to read best practices: {}", e))?;
        let bp_docs: Vec<KnowledgeDocument> = serde_json::from_str(&bp_content)?;

        // Load extension documents
        let ext_path = format!("{}/extension-docs.json", knowledge_base_dir);
        let ext_content = std::fs::read_to_string(&ext_path)
            .map_err(|e| anyhow::anyhow!("Failed to read extensions: {}", e))?;
        let ext_docs: Vec<KnowledgeDocument> = serde_json::from_str(&ext_content)?;

        // Load relationships
        let rel_path = format!("{}/relationships.json", knowledge_base_dir);
        let rel_content = std::fs::read_to_string(&rel_path)
            .map_err(|e| anyhow::anyhow!("Failed to read relationships: {}", e))?;
        let relationships: Vec<Relationship> = serde_json::from_str(&rel_content)?;

        // Build document index
        let mut documents = std::collections::HashMap::new();
        for doc in bp_docs.into_iter().chain(ext_docs.into_iter()) {
            documents.insert(doc.id.clone(), doc);
        }

        // Update knowledge base
        let mut kb = self.knowledge_base.write().await;
        kb.documents = documents;
        kb.relationships = relationships;

        info!(
            "Knowledge base loaded: {} documents, {} relationships",
            kb.documents.len(),
            kb.relationships.len()
        );

        Ok(())
    }

    /// Get service address
    pub fn addr(&self) -> SocketAddr {
        self.addr
    }

    /// Search knowledge base by keyword and category
    async fn search_knowledge(
        &self,
        query: &str,
        category: Option<&str>,
    ) -> Vec<KnowledgeDocument> {
        let kb = self.knowledge_base.read().await;
        let query_lower = query.to_lowercase();

        kb.documents
            .values()
            .filter(|doc| {
                let matches_query = doc.content.to_lowercase().contains(&query_lower)
                    || doc
                        .tags
                        .iter()
                        .any(|t| t.to_lowercase().contains(&query_lower));

                let matches_category = category.map(|c| doc.category == c).unwrap_or(true);

                matches_query && matches_category
            })
            .cloned()
            .collect()
    }

    /// Call an MCP tool via registry
    pub async fn call_mcp_tool(&self, req: McpToolRequest) -> Result<McpToolResponse> {
        let start = std::time::Instant::now();

        debug!("Calling MCP tool: {}", req.tool_name);

        let result = self.execute_mcp_tool(&req.tool_name, &req.args).await?;
        let duration_ms = start.elapsed().as_millis() as u64;

        Ok(McpToolResponse {
            result,
            duration_ms,
        })
    }

    /// Execute MCP tool with arguments via registry
    async fn execute_mcp_tool(
        &self,
        tool_name: &str,
        args: &serde_json::Value,
    ) -> Result<serde_json::Value> {
        // Check if tool exists in registry
        if !self.tool_registry.has_tool(tool_name) {
            return Err(anyhow::anyhow!("Unknown tool: {}", tool_name));
        }

        // Execute tool through registry
        match self.tool_registry.execute(tool_name, args).await {
            Ok(result) => Ok(result),
            Err(e) => Err(anyhow::anyhow!("Tool execution failed: {}", e)),
        }
    }

    /// Execute suggested tools and collect results
    async fn execute_tools(
        &self,
        suggestions: &[crate::tool_integration::ToolSuggestion],
        max_tools: usize,
    ) -> Option<(
        Vec<crate::mcp::ToolExecution>,
        Vec<(String, serde_json::Value)>,
    )> {
        if suggestions.is_empty() {
            return None;
        }

        debug!(
            "Executing {} suggested tools in hybrid mode",
            suggestions.len().min(max_tools)
        );

        let mut executions = Vec::new();
        let mut results = Vec::new();

        for suggestion in suggestions.iter().take(max_tools) {
            match self
                .tool_registry
                .execute(&suggestion.tool_name, &suggestion.args)
                .await
            {
                Ok(result) => {
                    debug!("Tool {} executed successfully", suggestion.tool_name);
                    results.push((suggestion.tool_name.clone(), result.clone()));
                    executions.push(crate::mcp::ToolExecution {
                        tool_name: suggestion.tool_name.clone(),
                        result: result.clone(),
                        duration_ms: 0,
                    });
                }
                Err(e) => {
                    debug!("Tool {} execution failed: {}", suggestion.tool_name, e);
                }
            }
        }

        if executions.is_empty() {
            None
        } else {
            Some((executions, results))
        }
    }

    /// Ask AI a question using RAG with optional hybrid tool execution
    pub async fn ask(&self, req: AskRequest) -> Result<AskResponse> {
        debug!("Processing RAG question: {}", req.question);

        // Search knowledge base for relevant documents
        let search_results = self.search_knowledge(&req.question, None).await;

        if search_results.is_empty() {
            return Ok(AskResponse {
                answer: "I couldn't find any relevant information in the knowledge base for this \
                         question."
                    .to_string(),
                sources: vec![],
                confidence: 20,
                reasoning: "No matching documents found".to_string(),
                tool_executions: None,
            });
        }

        // Sort by relevance (best practices have explicit relevance scores)
        let mut results = search_results;
        results.sort_by(|a, b| {
            let a_rel = a.relevance.unwrap_or(50);
            let b_rel = b.relevance.unwrap_or(50);
            b_rel.cmp(&a_rel)
        });

        // Get top 3 most relevant documents
        let top_results: Vec<_> = results.iter().take(3).collect();

        // Build answer from top results
        let mut answer_parts =
            vec!["Based on the knowledge base, here's what I found:".to_string()];

        for doc in &top_results {
            if let Some(title) = &doc.title {
                answer_parts.push(format!(
                    "\n- **{}**: {}",
                    title,
                    &doc.content[..std::cmp::min(150, doc.content.len())]
                ));
            } else if let Some(name) = &doc.name {
                answer_parts.push(format!(
                    "\n- **{}**: {}",
                    name,
                    &doc.content[..std::cmp::min(150, doc.content.len())]
                ));
            }
        }

        let mut answer = answer_parts.join("\n");
        let sources: Vec<String> = top_results
            .iter()
            .filter_map(|d| d.title.clone().or_else(|| d.name.clone()))
            .collect();

        let confidence = (top_results.iter().filter_map(|d| d.relevance).sum::<u8>() as f32
            / top_results.len() as f32) as u8;

        // Handle hybrid execution mode (auto-trigger tools if enabled)
        let mut tool_executions = None;
        if req.enable_tool_execution.unwrap_or(false) {
            let max_tools = req.max_tool_calls.unwrap_or(3) as usize;
            let tool_suggestions =
                crate::tool_integration::analyze_for_tools(&answer, &req.question);

            if let Some((executions, results)) =
                self.execute_tools(&tool_suggestions, max_tools).await
            {
                if !results.is_empty() {
                    answer = crate::tool_integration::enrich_answer_with_results(answer, &results);
                    tool_executions = Some(executions);
                }
            }
        }

        Ok(AskResponse {
            answer,
            sources,
            confidence,
            reasoning: format!(
                "Retrieved {} relevant documents using keyword search across {} total documents",
                top_results.len(),
                results.len()
            ),
            tool_executions,
        })
    }

    /// Get extension dependency DAG
    pub async fn get_extension_dag(&self) -> Result<DagResponse> {
        debug!("Building extension DAG");

        let kb = self.knowledge_base.read().await;

        // Build nodes from extension documents
        let nodes: Vec<DagNode> = kb
            .documents
            .values()
            .filter(|doc| doc.doc_type == "extension_metadata")
            .map(|doc| DagNode {
                name: doc.name.clone().unwrap_or_else(|| doc.id.clone()),
                dependencies: doc.dependencies.clone().unwrap_or_default(),
                version: "1.0.0".to_string(),
            })
            .collect();

        // Build edges from dependency relationships
        let edges: Vec<(String, String)> = kb
            .relationships
            .iter()
            .filter(|rel| rel.relationship_type == "depends_on")
            .map(|rel| {
                let source = rel.source_id.strip_prefix("ext_").unwrap_or(&rel.source_id);
                let target = rel.target_id.strip_prefix("ext_").unwrap_or(&rel.target_id);
                (source.to_string(), target.to_string())
            })
            .collect();

        Ok(DagResponse { nodes, edges })
    }

    /// Get best practices for a category
    pub async fn get_best_practices(&self, category: &str) -> Result<Vec<BestPractice>> {
        debug!("Retrieving best practices for category: {}", category);

        let kb = self.knowledge_base.read().await;

        // Filter documents by category and type
        let mut practices: Vec<BestPractice> = kb
            .documents
            .values()
            .filter(|doc| doc.category == category && doc.doc_type == "best_practice")
            .map(|doc| BestPractice {
                title: doc.title.clone().unwrap_or_else(|| doc.id.clone()),
                description: doc.content.clone(),
                category: doc.category.clone(),
                relevance: doc.relevance.unwrap_or(70),
            })
            .collect();

        // Sort by relevance descending
        practices.sort_by(|a, b| b.relevance.cmp(&a.relevance));

        Ok(practices)
    }

    /// Health check endpoint
    pub async fn health_check(&self) -> Result<()> {
        Ok(())
    }

    /// Get all available tools
    pub fn list_all_tools(&self) -> Vec<crate::mcp::ToolDefinition> {
        self.tool_registry.list_tools()
    }

    /// Get tools by category
    pub fn tools_by_category(
        &self,
        category: crate::mcp::ToolCategory,
    ) -> Vec<crate::mcp::ToolDefinition> {
        self.tool_registry.tools_by_category(category)
    }

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

impl Default for AiService {
    fn default() -> Self {
        Self::new("127.0.0.1:8083".parse().unwrap())
    }
}