velocia 0.3.5

//! `RemoteAgentConnections` – manages HTTP connections to remote A2A agents.
//!
//! Mirrors Python's `RemoteAgentConnections` class.

use std::collections::HashMap;
use std::sync::Arc;

use serde_json::Value;
use tracing::{error, info};
use uuid::Uuid;

use crate::a2a::client::send_a2a_message;
use crate::a2a::types::{AgentCard, Message, Part, Task};
use crate::config::auth::{AuthConfig, AuthType};
use crate::error::{AgentKitError, Result};

// ── Remote connection ─────────────────────────────────────────────────────────

#[derive(Clone)]
pub struct RemoteAgentConnection {
    pub card: AgentCard,
    pub auth: AuthConfig,
    http: reqwest::Client,
}

impl RemoteAgentConnection {
    pub fn new(card: AgentCard, auth: AuthConfig) -> Self {
        Self {
            card,
            auth,
            http: reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(60))
                .build()
                .expect("reqwest client build failed"),
        }
    }

    /// Send a message to the remote A2A agent via JSON-RPC. Uses `message/stream`
    /// (SSE) when the agent's card advertises streaming support, `message/send`
    /// (single JSON response) otherwise — matching the reference `a2a-sdk` client.
    pub async fn send_message(
        &self,
        request: Message,
        propagation_headers: HashMap<String, String>,
    ) -> Result<Task> {
        info!("A2A → {} ({})", self.card.name, self.card.url);

        if self.auth.auth_type != AuthType::NoAuth {
            // TODO: plug in CognitoM2MCredentialService::get_credentials() here
        }

        let task = send_a2a_message(&self.http, &self.card, request, propagation_headers).await?;
        info!("A2A ← {} — {} artifact(s) collected", self.card.name, task.artifacts.len());

        Ok(task)
    }
}

// ── Multi-agent manager ───────────────────────────────────────────────────────

/// Manages connections to all configured remote agents.
pub struct RemoteAgentManager {
    connections: HashMap<String, Arc<RemoteAgentConnection>>,
    pub agents_description: String,
}

impl RemoteAgentManager {
    pub fn new() -> Self {
        Self {
            connections: HashMap::new(),
            agents_description: String::new(),
        }
    }

    /// Discover and connect to remote agents from their addresses.
    pub async fn connect(&mut self, addresses: Vec<(String, AuthConfig)>) -> Result<()> {
        let http = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(60))
            .build()
            .map_err(|e| AgentKitError::A2aClient(e.to_string()))?;

        for (i, (url, auth)) in addresses.into_iter().enumerate() {
            let card_url = format!("{url}/.well-known/agent.json");
            info!("Connecting to remote agent {i}: {url}");

            match http.get(&card_url).send().await {
                Ok(resp) => match resp.json::<AgentCard>().await {
                    Ok(card) => {
                        info!("Connected to remote agent '{}'", card.name);
                        self.connections.insert(
                            card.name.clone(),
                            Arc::new(RemoteAgentConnection::new(card, auth)),
                        );
                    }
                    Err(e) => error!("Failed to parse agent card from {url}: {e}"),
                },
                Err(e) => error!("Failed to reach remote agent at {url}: {e}"),
            }
        }

        self.agents_description = self
            .connections
            .values()
            .map(|c| {
                serde_json::json!({
                    "name": c.card.name,
                    "description": c.card.description
                })
                .to_string()
            })
            .collect::<Vec<_>>()
            .join("\n");

        info!(
            "Remote agent manager ready with {} agent(s)",
            self.connections.len()
        );
        Ok(())
    }

    /// Returns all established connections as `Arc` — used to wire sub-agents.
    pub fn take_connections(&self) -> Vec<Arc<RemoteAgentConnection>> {
        self.connections.values().cloned().collect()
    }

    pub fn list_agents(&self) -> Vec<Value> {
        self.connections
            .values()
            .map(|c| serde_json::json!({"name": c.card.name, "description": c.card.description}))
            .collect()
    }

    /// Delegate a task to a named remote agent.
    pub async fn send_message(
        &self,
        agent_name: &str,
        task: &str,
        context_id: Option<String>,
        task_id: Option<String>,
        propagation_headers: HashMap<String, String>,
    ) -> Result<Task> {
        let conn = self.connections.get(agent_name).ok_or_else(|| {
            AgentKitError::RemoteAgentNotFound { name: agent_name.to_string() }
        })?;

        let ctx_id = context_id.unwrap_or_else(|| Uuid::new_v4().to_string());

        let message = Message {
            kind: "message".to_string(),
            role: crate::a2a::types::Role::User,
            parts: vec![Part::Text { text: task.to_string(), metadata: None }],
            message_id: Uuid::new_v4().to_string(),
            context_id: Some(ctx_id),
            task_id,
            metadata: None,
        };

        conn.send_message(message, propagation_headers).await
    }
}

impl Default for RemoteAgentManager {
    fn default() -> Self { Self::new() }
}