Vapora/crates/vapora-agents/src/swarm_adapter.rs

use std::sync::Arc;

use async_trait::async_trait;
use uuid::Uuid;
use vapora_swarm::coordinator::SwarmCoordinator;

use crate::coordination::{AgentAssignment, AgentLoad, AgentProfile, SwarmCoordination};

/// Adapter: SwarmCoordination → SwarmCoordinator
/// Implements the coordination abstraction using the real swarm coordinator.
pub struct SwarmCoordinationAdapter {
    swarm: Arc<SwarmCoordinator>,
}

impl SwarmCoordinationAdapter {
    pub fn new(swarm: Arc<SwarmCoordinator>) -> Self {
        Self { swarm }
    }
}

#[async_trait]
impl SwarmCoordination for SwarmCoordinationAdapter {
    async fn register_profiles(&self, profiles: Vec<AgentProfile>) -> anyhow::Result<()> {
        for profile in profiles {
            let swarm_profile = vapora_swarm::messages::AgentProfile {
                id: profile.id.clone(),
                roles: vec![profile.role.clone()],
                capabilities: vec![profile.role],
                current_load: 0.0,
                availability: true,
                success_rate: profile.success_rate,
            };
            self.swarm.register_agent(swarm_profile)?;
        }
        Ok(())
    }

    /// Select best agent via swarm bidding.
    ///
    /// Uses `submit_task_for_bidding` which applies load-balanced scoring
    /// (success_rate / (1 + current_load)) across all available agents with
    /// matching capabilities.
    async fn select_agent(
        &self,
        task_type: &str,
        required_expertise: Option<&str>,
    ) -> anyhow::Result<AgentAssignment> {
        let capabilities: Vec<String> = match required_expertise {
            Some(exp) => vec![task_type.to_string(), exp.to_string()],
            None => vec![task_type.to_string()],
        };

        // Use a ephemeral task_id for selection; the caller manages actual task IDs.
        let selection_id = Uuid::new_v4().to_string();

        let agent_id = self
            .swarm
            .submit_task_for_bidding(selection_id, task_type.to_string(), capabilities)
            .await
            .map_err(|e| anyhow::anyhow!("Swarm bidding failed: {}", e))?
            .ok_or_else(|| anyhow::anyhow!("No available agent for task_type: {}", task_type))?;

        let confidence = self
            .swarm
            .get_agent(&agent_id)
            .map(|profile| profile.success_rate)
            .unwrap_or(0.5);

        Ok(AgentAssignment {
            // Swarm profiles use ID as display name (no separate name field)
            agent_name: agent_id.clone(),
            agent_id,
            confidence,
        })
    }

    /// Report task completion and update agent load in the swarm.
    ///
    /// On success the agent is marked available with minimal load.
    /// On failure the agent receives a penalty load (0.5) to deprioritize it
    /// in future selections until it recovers.
    async fn report_completion(
        &self,
        agent_id: &str,
        success: bool,
        _duration_ms: u64,
    ) -> anyhow::Result<()> {
        let new_load = if success { 0.0 } else { 0.5 };
        self.swarm
            .update_agent_status(agent_id, new_load, true)
            .map_err(|e| anyhow::anyhow!("Failed to update agent status: {}", e))
    }

    /// Query current agent load from swarm profile.
    ///
    /// Infers `current_tasks` from the fractional load stored in the swarm
    /// profile (each task represents ~10% of a capacity-10 agent).
    async fn agent_load(&self, agent_id: &str) -> anyhow::Result<AgentLoad> {
        let profile = self
            .swarm
            .get_agent(agent_id)
            .ok_or_else(|| anyhow::anyhow!("Agent not found in swarm: {}", agent_id))?;

        const CAPACITY: usize = 10;
        let current_tasks = (profile.current_load * CAPACITY as f64).round() as usize;

        Ok(AgentLoad {
            agent_id: agent_id.to_string(),
            current_tasks,
            capacity: CAPACITY,
        })
    }
}