Vapora/crates/vapora-backend/src/services/agent_service.rs

// Agent service - Registry and management for the 12 agent roles

use chrono::Utc;
use surrealdb::engine::remote::ws::Client;
use surrealdb::Surreal;
use vapora_shared::models::{Agent, AgentRole, AgentStatus};
use vapora_shared::{Result, VaporaError};

/// Service for managing agents
#[derive(Clone)]
pub struct AgentService {
    db: Surreal<Client>,
}

impl AgentService {
    /// Create a new AgentService instance
    pub fn new(db: Surreal<Client>) -> Self {
        Self { db }
    }

    /// Register a new agent
    pub async fn register_agent(&self, mut agent: Agent) -> Result<Agent> {
        // Set creation timestamp
        agent.created_at = Utc::now();

        // Check if agent with this role already exists
        let existing = self.get_agent_by_role(&agent.role).await;
        if existing.is_ok() {
            return Err(VaporaError::InvalidInput(format!(
                "Agent with role '{:?}' already exists",
                agent.role
            )));
        }

        // Create agent in database
        let created: Option<Agent> = self
            .db
            .create("agents")
            .content(agent)
            .await?
            .into_iter()
            .next();

        created.ok_or_else(|| VaporaError::DatabaseError("Failed to register agent".to_string()))
    }

    /// List all agents
    pub async fn list_agents(&self) -> Result<Vec<Agent>> {
        let mut response = self
            .db
            .query("SELECT * FROM agents ORDER BY role ASC")
            .await?;

        let agents: Vec<Agent> = response.take(0)?;
        Ok(agents)
    }

    /// List agents by status
    pub async fn list_agents_by_status(&self, status: AgentStatus) -> Result<Vec<Agent>> {
        let status_str = match status {
            AgentStatus::Active => "active",
            AgentStatus::Inactive => "inactive",
            AgentStatus::Updating => "updating",
            AgentStatus::Error => "error",
        };

        let mut response = self
            .db
            .query("SELECT * FROM agents WHERE status = $status ORDER BY role ASC")
            .bind(("status", status_str.to_string()))
            .await?;

        let agents: Vec<Agent> = response.take(0)?;
        Ok(agents)
    }

    /// Get an agent by ID
    pub async fn get_agent(&self, id: &str) -> Result<Agent> {
        let agent: Option<Agent> = self.db.select(("agents", id)).await?;

        agent.ok_or_else(|| VaporaError::NotFound(format!("Agent with id '{}' not found", id)))
    }

    /// Get an agent by role
    pub async fn get_agent_by_role(&self, role: &AgentRole) -> Result<Agent> {
        let role_str = match role {
            AgentRole::Architect => "architect",
            AgentRole::Developer => "developer",
            AgentRole::CodeReviewer => "code_reviewer",
            AgentRole::Tester => "tester",
            AgentRole::Documenter => "documenter",
            AgentRole::Marketer => "marketer",
            AgentRole::Presenter => "presenter",
            AgentRole::DevOps => "dev_ops",
            AgentRole::Monitor => "monitor",
            AgentRole::Security => "security",
            AgentRole::ProjectManager => "project_manager",
            AgentRole::DecisionMaker => "decision_maker",
        };

        let mut response = self
            .db
            .query("SELECT * FROM agents WHERE role = $role LIMIT 1")
            .bind(("role", role_str.to_string()))
            .await?;

        let agents: Vec<Agent> = response.take(0)?;

        agents
            .into_iter()
            .next()
            .ok_or_else(|| VaporaError::NotFound(format!("Agent with role '{:?}' not found", role)))
    }

    /// Update an agent
    pub async fn update_agent(&self, id: &str, mut updates: Agent) -> Result<Agent> {
        // Verify agent exists
        let existing = self.get_agent(id).await?;

        // Preserve certain fields
        updates.id = existing.id;
        updates.created_at = existing.created_at;

        // Update in database
        let updated: Option<Agent> = self.db.update(("agents", id)).content(updates).await?;

        updated.ok_or_else(|| VaporaError::DatabaseError("Failed to update agent".to_string()))
    }

    /// Update agent status
    pub async fn update_agent_status(&self, id: &str, status: AgentStatus) -> Result<Agent> {
        // Verify agent exists
        self.get_agent(id).await?;

        let updated: Option<Agent> = self
            .db
            .update(("agents", id))
            .merge(serde_json::json!({
                "status": status
            }))
            .await?;

        updated
            .ok_or_else(|| VaporaError::DatabaseError("Failed to update agent status".to_string()))
    }

    /// Add capability to an agent
    pub async fn add_capability(&self, id: &str, capability: String) -> Result<Agent> {
        let mut agent = self.get_agent(id).await?;

        // Add capability if not already present
        if !agent.capabilities.contains(&capability) {
            agent.capabilities.push(capability);

            let updated: Option<Agent> = self
                .db
                .update(("agents", id))
                .merge(serde_json::json!({
                    "capabilities": agent.capabilities
                }))
                .await?;

            return updated
                .ok_or_else(|| VaporaError::DatabaseError("Failed to add capability".to_string()));
        }

        Ok(agent)
    }

    /// Remove capability from an agent
    pub async fn remove_capability(&self, id: &str, capability: &str) -> Result<Agent> {
        let mut agent = self.get_agent(id).await?;

        // Remove capability
        agent.capabilities.retain(|c| c != capability);

        let updated: Option<Agent> = self
            .db
            .update(("agents", id))
            .merge(serde_json::json!({
                "capabilities": agent.capabilities
            }))
            .await?;

        updated.ok_or_else(|| VaporaError::DatabaseError("Failed to remove capability".to_string()))
    }

    /// Add skill to an agent
    pub async fn add_skill(&self, id: &str, skill: String) -> Result<Agent> {
        let mut agent = self.get_agent(id).await?;

        // Add skill if not already present
        if !agent.skills.contains(&skill) {
            agent.skills.push(skill);

            let updated: Option<Agent> = self
                .db
                .update(("agents", id))
                .merge(serde_json::json!({
                    "skills": agent.skills
                }))
                .await?;

            return updated
                .ok_or_else(|| VaporaError::DatabaseError("Failed to add skill".to_string()));
        }

        Ok(agent)
    }

    /// Deregister an agent
    pub async fn deregister_agent(&self, id: &str) -> Result<()> {
        // Verify agent exists
        self.get_agent(id).await?;

        // Delete from database
        let _: Option<Agent> = self.db.delete(("agents", id)).await?;

        Ok(())
    }

    /// Get agent health status (checks if agent is active and responding)
    pub async fn check_agent_health(&self, id: &str) -> Result<bool> {
        let agent = self.get_agent(id).await?;
        Ok(agent.status == AgentStatus::Active)
    }

    /// Get agents available for task assignment (active agents with capacity)
    pub async fn get_available_agents(&self) -> Result<Vec<Agent>> {
        let mut response = self
            .db
            .query("SELECT * FROM agents WHERE status = 'active' ORDER BY role ASC")
            .await?;

        let agents: Vec<Agent> = response.take(0)?;
        Ok(agents)
    }
}

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

    // Note: These are placeholder tests. Real tests require a running SurrealDB instance
    // or mocking. For Phase 1, we'll add integration tests that use a test database.

    #[test]
    fn test_agent_service_creation() {
        // This test just verifies the service can be created
        // Real database tests will be in integration tests
    }
}