swarm-gateway 0.1.0

use axum::extract::State;
use axum::http::StatusCode;
use axum::response::Json;
use serde::{Deserialize, Serialize};
use swarm_core::Agent;
use swarm_core::agent::AgentId;
use swarm_core::message::Message;

use crate::state::AppState;

#[derive(Serialize)]
pub struct HealthResponse {
    pub status: &'static str,
    pub name: String,
    pub version: &'static str,
    pub agents: usize,
}

pub async fn health(State(state): State<AppState>) -> Json<HealthResponse> {
    let agents = state.orchestrator.agents().read().await;
    Json(HealthResponse {
        status: "ok",
        name: state.orchestrator.config().name.clone(),
        version: env!("CARGO_PKG_VERSION"),
        agents: agents.len(),
    })
}

#[derive(Serialize)]
pub struct AgentInfo {
    pub id: String,
    pub name: String,
    pub role: String,
    pub status: String,
    pub model: String,
    pub provider: String,
}

pub async fn list_agents(State(state): State<AppState>) -> Json<Vec<AgentInfo>> {
    let agents = state.orchestrator.agents().read().await;
    let infos = agents
        .values()
        .map(|h| AgentInfo {
            id: h.id.to_string(),
            name: h.manifest.name.clone(),
            role: format!("{:?}", h.manifest.role),
            status: format!("{:?}", h.status),
            model: h.manifest.model.clone(),
            provider: h.manifest.llm_provider.clone(),
        })
        .collect();
    Json(infos)
}

#[derive(Deserialize)]
pub struct ChatRequest {
    pub agent_id: String,
    pub message: String,
}

#[derive(Serialize)]
pub struct ChatResponse {
    pub agent_id: String,
    pub response: String,
}

pub async fn chat(
    State(state): State<AppState>,
    Json(req): Json<ChatRequest>,
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
    let agent_id = AgentId::new(&req.agent_id);

    // Get the provider for this agent
    let agents = state.orchestrator.agents().read().await;
    let handle = agents.get(&agent_id).ok_or_else(|| {
        (
            StatusCode::NOT_FOUND,
            format!("Agent not found: {}", req.agent_id),
        )
    })?;

    let provider_name = handle.manifest.llm_provider.clone();
    let manifest = handle.manifest.clone();
    drop(agents);

    let provider = state.providers.get(&provider_name).ok_or_else(|| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Provider not found: {provider_name}"),
        )
    })?;

    // Create a temporary agent to process the message
    let mut agent =
        swarm_agents::BaseAgent::new(manifest, provider.clone(), state.tools.clone());

    let input = Message::user(AgentId::new("user"), &req.message);
    let response = agent.process(input).await.map_err(|e: swarm_core::SwarmError| {
        (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    Ok(Json(ChatResponse {
        agent_id: req.agent_id,
        response: response.content,
    }))
}

#[derive(Serialize)]
pub struct ProvidersResponse {
    pub providers: Vec<String>,
}

pub async fn list_providers(State(state): State<AppState>) -> Json<ProvidersResponse> {
    Json(ProvidersResponse {
        providers: state.providers.list().into_iter().map(String::from).collect(),
    })
}

#[derive(Serialize)]
pub struct ToolInfo {
    pub name: String,
}

pub async fn list_tools(State(state): State<AppState>) -> Json<Vec<ToolInfo>> {
    Json(
        state
            .tools
            .list()
            .into_iter()
            .map(|name| ToolInfo {
                name: name.to_string(),
            })
            .collect(),
    )
}

#[derive(Deserialize)]
pub struct SendMessageRequest {
    pub from: String,
    pub to: String,
    pub message: String,
}

pub async fn send_message(
    State(state): State<AppState>,
    Json(req): Json<SendMessageRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
    let from = AgentId::new(&req.from);
    let to = AgentId::new(&req.to);
    let msg = Message::agent_to_agent(from, to.clone(), &req.message);

    state
        .orchestrator
        .send_to_agent(&to, msg)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(serde_json::json!({ "status": "sent" })))
}

// ---------------------------------------------------------------------------
// Apple FM Agent Router — on-device task routing via Apple Intelligence
// ---------------------------------------------------------------------------

#[derive(Deserialize)]
pub struct RouteRequest {
    pub task: String,
    #[serde(default)]
    pub context: String,
    #[serde(default)]
    pub auto_forward: bool,
}

pub async fn route_task(
    State(state): State<AppState>,
    Json(req): Json<RouteRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
    // Build agent descriptions from live registry
    let agents = state.orchestrator.agents().read().await;
    let agent_descs: Vec<String> = agents
        .values()
        .map(|h| {
            format!(
                "{} — {}",
                h.manifest.name,
                h.manifest.system_prompt.chars().take(100).collect::<String>()
            )
        })
        .collect();
    drop(agents);

    // Call Apple FM agent_router via CLI (on-device, zero cost)
    let agents_str = agent_descs.join("\n");
    let mut cmd = tokio::process::Command::new("apple-fm");
    cmd.args([
        "agent_router",
        "--task",
        &req.task,
        "--context",
        &format!("{}\n\nAvailable agents:\n{}", req.context, agents_str),
        "--quiet",
    ]);

    let output = cmd.output().await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Apple FM not available: {e}"),
        )
    })?;

    if !output.status.success() {
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            "Apple FM routing failed".to_string(),
        ));
    }

    let raw = String::from_utf8_lossy(&output.stdout);

    // Parse the routing decision
    let routing: serde_json::Value = serde_json::from_str(raw.trim())
        .or_else(|_| {
            // Apple FM sometimes wraps in ```json ... ```
            let stripped = raw
                .trim()
                .strip_prefix("```json")
                .unwrap_or(raw.trim())
                .strip_suffix("```")
                .unwrap_or(raw.trim());
            serde_json::from_str(stripped)
        })
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to parse routing: {e}\nRaw: {raw}"),
            )
        })?;

    // If auto_forward is true, send the task to the primary agent
    if req.auto_forward {
        if let Some(primary) = routing.get("primary").and_then(|v| v.as_str()) {
            let agents = state.orchestrator.agents().read().await;
            // Find agent by name (case-insensitive partial match)
            let matched = agents.values().find(|h| {
                h.manifest
                    .name
                    .to_lowercase()
                    .contains(&primary.to_lowercase())
            });

            if let Some(handle) = matched {
                let agent_id = handle.id.clone();
                let manifest = handle.manifest.clone();
                let provider_name = handle.manifest.llm_provider.clone();
                drop(agents);

                if let Some(provider) = state.providers.get(&provider_name) {
                    let mut agent = swarm_agents::BaseAgent::new(
                        manifest,
                        provider.clone(),
                        state.tools.clone(),
                    );
                    let input = Message::user(AgentId::new("user"), &req.task);
                    match agent.process(input).await {
                        Ok(response) => {
                            return Ok(Json(serde_json::json!({
                                "routing": routing,
                                "forwarded": true,
                                "agent_id": agent_id.to_string(),
                                "response": response.content,
                            })));
                        }
                        Err(e) => {
                            return Ok(Json(serde_json::json!({
                                "routing": routing,
                                "forwarded": false,
                                "forward_error": e.to_string(),
                            })));
                        }
                    }
                }
            }
        }
    }

    Ok(Json(serde_json::json!({
        "routing": routing,
        "forwarded": false,
    })))
}