prvng_platform/crates/rag/src/context.rs

//! Workspace context extraction and management

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

/// Workspace context for RAG enrichment
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkspaceContext {
    /// Active workspace name
    pub workspace_name: String,

    /// Workspace path
    pub workspace_path: String,

    /// Active infrastructure
    pub active_infrastructure: Option<String>,

    /// Available infrastructures and their status
    pub infrastructures: Vec<InfrastructureState>,

    /// Installed taskservs
    pub installed_taskservs: Vec<String>,

    /// Installed providers
    pub installed_providers: Vec<String>,

    /// Recent deployments (last N)
    pub recent_deployments: Vec<DeploymentInfo>,

    /// Capabilities of the workspace
    pub capabilities: CapabilitySet,

    /// Custom metadata
    pub metadata: HashMap<String, String>,
}

/// Infrastructure state within a workspace
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InfrastructureState {
    /// Infrastructure name
    pub name: String,

    /// Provider (upcloud, aws, local)
    pub provider: String,

    /// Number of servers
    pub server_count: usize,

    /// Deployment phase (New, Configured, Deploying, Running, Stopping, Error)
    pub deployment_phase: String,

    /// Deployed taskservs
    pub taskservs: Vec<String>,

    /// Deployed clusters
    pub clusters: Vec<String>,

    /// Last updated timestamp
    pub last_updated: Option<String>,
}

/// Deployment information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeploymentInfo {
    /// Unique deployment ID
    pub id: String,

    /// Infrastructure name
    pub infrastructure: String,

    /// Deployment type (server_create, taskserv_install, cluster_deploy)
    pub event_type: String,

    /// Status (success, failed, partial)
    pub status: String,

    /// Timestamp
    pub timestamp: String,

    /// Duration in seconds
    pub duration_seconds: Option<u64>,

    /// Error message (if failed)
    pub error_message: Option<String>,
}

/// Capability set for workspace
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CapabilitySet {
    /// Can deploy to these providers
    pub deployable_providers: Vec<String>,

    /// Can create these resource types
    pub creatable_resources: Vec<String>,

    /// Installed tools
    pub installed_tools: Vec<String>,

    /// Supported features
    pub features: Vec<String>,
}

impl WorkspaceContext {
    /// Create a new workspace context
    pub fn new(workspace_name: String, workspace_path: String) -> Self {
        Self {
            workspace_name,
            workspace_path,
            active_infrastructure: None,
            infrastructures: Vec::new(),
            installed_taskservs: Vec::new(),
            installed_providers: Vec::new(),
            recent_deployments: Vec::new(),
            capabilities: CapabilitySet::default(),
            metadata: HashMap::new(),
        }
    }

    /// Add infrastructure state
    pub fn add_infrastructure(&mut self, infra: InfrastructureState) {
        self.infrastructures.push(infra);
    }

    /// Add taskserv to installed list
    pub fn add_taskserv(&mut self, name: String) {
        if !self.installed_taskservs.contains(&name) {
            self.installed_taskservs.push(name);
        }
    }

    /// Add provider to installed list
    pub fn add_provider(&mut self, name: String) {
        if !self.installed_providers.contains(&name) {
            self.installed_providers.push(name);
        }
    }

    /// Add deployment to recent list
    pub fn add_deployment(&mut self, deployment: DeploymentInfo) {
        self.recent_deployments.insert(0, deployment);
        // Keep only last 10 deployments
        if self.recent_deployments.len() > 10 {
            self.recent_deployments.truncate(10);
        }
    }

    /// Get context summary as human-readable string
    pub fn summarize(&self) -> String {
        let mut summary = format!(
            "Workspace: {} (Path: {})\n",
            self.workspace_name, self.workspace_path
        );

        if let Some(active) = &self.active_infrastructure {
            summary.push_str(&format!("Active Infrastructure: {}\n", active));
        }

        if !self.installed_taskservs.is_empty() {
            summary.push_str(&format!(
                "Taskservs: {}\n",
                self.installed_taskservs.join(", ")
            ));
        }

        if !self.installed_providers.is_empty() {
            summary.push_str(&format!(
                "Providers: {}\n",
                self.installed_providers.join(", ")
            ));
        }

        if !self.recent_deployments.is_empty() {
            let recent = &self.recent_deployments[0];
            summary.push_str(&format!(
                "Latest deployment: {} ({}) - {}\n",
                recent.event_type, recent.infrastructure, recent.status
            ));
        }

        summary
    }

    /// Enrich a query with workspace context
    pub fn enrich_query(&self, query: &str) -> String {
        format!("{}\n\nWorkspace Context:\n{}\n", query, self.summarize())
    }
}

impl Default for CapabilitySet {
    fn default() -> Self {
        Self {
            deployable_providers: vec![],
            creatable_resources: vec![
                "servers".to_string(),
                "taskservs".to_string(),
                "clusters".to_string(),
            ],
            installed_tools: vec![],
            features: vec![],
        }
    }
}

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

    #[test]
    fn test_workspace_context_creation() {
        let ctx = WorkspaceContext::new(
            "test-workspace".to_string(),
            "/path/to/workspace".to_string(),
        );

        assert_eq!(ctx.workspace_name, "test-workspace");
        assert_eq!(ctx.workspace_path, "/path/to/workspace");
        assert!(ctx.infrastructures.is_empty());
    }

    #[test]
    fn test_add_taskserv() {
        let mut ctx = WorkspaceContext::new("test".to_string(), "/path".to_string());

        ctx.add_taskserv("kubernetes".to_string());
        ctx.add_taskserv("prometheus".to_string());

        assert_eq!(ctx.installed_taskservs.len(), 2);
        assert!(!ctx.installed_taskservs.contains(&"redis".to_string()));
    }

    #[test]
    fn test_enrich_query() {
        let mut ctx = WorkspaceContext::new(
            "librecloud".to_string(),
            "/workspace/librecloud".to_string(),
        );

        ctx.add_taskserv("kubernetes".to_string());

        let enriched = ctx.enrich_query("How do I add monitoring?");

        assert!(enriched.contains("How do I add monitoring?"));
        assert!(enriched.contains("Workspace Context"));
        assert!(enriched.contains("librecloud"));
        assert!(enriched.contains("kubernetes"));
    }
}