collet 0.1.0

//! REST API routes for swarm/hive agent graph visualization and control.

use std::sync::Arc;

use axum::Router;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::response::Json;
use axum::routing::{get, post};
use serde::{Deserialize, Serialize};

use super::state::AppState;

pub fn router() -> Router<Arc<AppState>> {
    Router::new()
        .route("/api/swarm/status", get(swarm_status))
        .route("/api/swarm/agent/{agent_id}/cancel", post(cancel_agent))
        .route("/api/swarm/agent/{agent_id}/extend", post(extend_agent))
}

// ── Types ────────────────────────────────────────────────────────────────

#[derive(Serialize, Clone)]
pub struct SwarmAgentNode {
    pub id: String,
    pub name: String,
    pub task: String,
    pub status: String, // "running" | "done" | "paused" | "error"
    pub success: bool,
    pub iteration: u32,
    pub tool_calls: u32,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub dependencies: Vec<String>,
    pub target_files: Vec<String>,
    pub modified_files: Vec<String>,
    /// Log entries for this agent (tool calls, responses, etc.)
    pub log: Vec<AgentLogEntry>,
}

#[derive(Serialize, Clone)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AgentLogEntry {
    ToolCall {
        name: String,
        args: String,
        timestamp: u64,
    },
    ToolResult {
        name: String,
        result: String,
        success: bool,
        timestamp: u64,
    },
    Token {
        text: String,
        timestamp: u64,
    },
    Response {
        text: String,
        timestamp: u64,
    },
    Status {
        message: String,
        timestamp: u64,
    },
}

#[derive(Serialize)]
pub struct SwarmGraphResponse {
    pub agents: Vec<SwarmAgentNode>,
    pub edges: Vec<(String, String)>, // (from_id, to_id) dependency edges
    pub phase: String,
    pub active: bool,
}

#[derive(Serialize)]
struct ControlResponse {
    ok: bool,
    message: String,
}

// ── Swarm Status ─────────────────────────────────────────────────────────

async fn swarm_status(State(state): State<Arc<AppState>>) -> Json<SwarmGraphResponse> {
    let agents = state.swarm_agents.read().await;
    let agent_nodes: Vec<SwarmAgentNode> = agents.values().cloned().collect();

    // Build edges from dependencies
    let mut edges = Vec::new();
    for node in &agent_nodes {
        for dep in &node.dependencies {
            if agents.contains_key(dep) {
                edges.push((dep.clone(), node.id.clone()));
            }
        }
    }

    let active = state.is_agent_active();

    Json(SwarmGraphResponse {
        agents: agent_nodes,
        edges,
        phase: String::new(),
        active,
    })
}

// ── Cancel ──────────────────────────────────────────────────────────────

async fn cancel_agent(
    State(state): State<Arc<AppState>>,
    Path(agent_id): Path<String>,
) -> (StatusCode, Json<ControlResponse>) {
    let knowledge = state.knowledge.read().await;
    if let Some(ref kb) = *knowledge {
        if kb.cancel_worker(&agent_id).await {
            (
                StatusCode::OK,
                Json(ControlResponse {
                    ok: true,
                    message: format!("Agent {agent_id} cancelled"),
                }),
            )
        } else {
            (
                StatusCode::NOT_FOUND,
                Json(ControlResponse {
                    ok: false,
                    message: format!("Agent {agent_id} not found or already finished"),
                }),
            )
        }
    } else {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ControlResponse {
                ok: false,
                message: "No active swarm session".into(),
            }),
        )
    }
}

// ── Extend ──────────────────────────────────────────────────────────────

#[derive(Deserialize)]
struct ExtendRequest {
    /// Number of additional iterations to grant.
    #[serde(default = "default_extend")]
    extra: u32,
}

fn default_extend() -> u32 {
    10
}

async fn extend_agent(
    State(state): State<Arc<AppState>>,
    Path(agent_id): Path<String>,
    Json(req): Json<ExtendRequest>,
) -> (StatusCode, Json<ControlResponse>) {
    let extra = if req.extra == 0 { 10 } else { req.extra };
    let knowledge = state.knowledge.read().await;
    if let Some(ref kb) = *knowledge {
        if kb.extend_worker(&agent_id, extra).await {
            (
                StatusCode::OK,
                Json(ControlResponse {
                    ok: true,
                    message: format!("Agent {agent_id} extended by {extra} iterations"),
                }),
            )
        } else {
            (
                StatusCode::NOT_FOUND,
                Json(ControlResponse {
                    ok: false,
                    message: format!("Agent {agent_id} not found or already finished"),
                }),
            )
        }
    } else {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ControlResponse {
                ok: false,
                message: "No active swarm session".into(),
            }),
        )
    }
}