axocoatl-server 0.0.1

use axum::{
    extract::{Path, State},
    http::StatusCode,
    Json,
};
use serde::{Deserialize, Serialize};

use crate::AppState;

// --- Health endpoints ---

#[derive(Serialize)]
pub struct HealthResponse {
    pub status: String,
    pub agents: usize,
}

pub async fn health(State(state): State<AppState>) -> Json<HealthResponse> {
    let daemon = state.read().await;
    Json(HealthResponse {
        status: "healthy".to_string(),
        agents: daemon.agent_count().await,
    })
}

pub async fn health_ready(State(state): State<AppState>) -> StatusCode {
    let daemon = state.read().await;
    if daemon.agent_count().await > 0 {
        StatusCode::OK
    } else {
        StatusCode::SERVICE_UNAVAILABLE
    }
}

pub async fn health_live() -> StatusCode {
    StatusCode::OK
}

// --- Agent endpoints ---

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

pub async fn list_agents(State(state): State<AppState>) -> Json<Vec<AgentInfo>> {
    let daemon = state.read().await;
    let agents: Vec<AgentInfo> = daemon
        .config
        .agents
        .iter()
        .map(|a| AgentInfo {
            id: a.id.clone(),
            name: a.name.clone(),
            provider: a.provider.clone(),
            model: a.model.clone(),
        })
        .collect();
    Json(agents)
}

#[derive(Deserialize)]
pub struct ExecuteRequest {
    pub input: String,
}

#[derive(Serialize)]
pub struct ExecuteResponse {
    pub output: String,
}

#[derive(Serialize)]
pub struct ErrorResponse {
    pub error: String,
}

pub async fn execute_agent(
    State(state): State<AppState>,
    Path(agent_id): Path<String>,
    Json(body): Json<ExecuteRequest>,
) -> Result<Json<ExecuteResponse>, (StatusCode, Json<ErrorResponse>)> {
    let daemon = state.read().await;
    match daemon.execute_agent(&agent_id, &body.input).await {
        Ok(output) => Ok(Json(ExecuteResponse { output: output.content })),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse {
                error: e.to_string(),
            }),
        )),
    }
}

#[derive(Serialize)]
pub struct AgentStatusResponse {
    pub agent_id: String,
    pub status: String,
}

pub async fn agent_status(
    State(state): State<AppState>,
    Path(agent_id): Path<String>,
) -> Result<Json<AgentStatusResponse>, (StatusCode, Json<ErrorResponse>)> {
    let daemon = state.read().await;
    let id = axocoatl_core::AgentId::new(&agent_id);

    match daemon.agent_registry.get(&id).await {
        Some(actor) => {
            let status = axocoatl_actor::get_agent_status(&actor)
                .await
                .unwrap_or_else(|e| axocoatl_core::AgentStatus::Failed {
                    error: e,
                    restarts: 0,
                });
            Ok(Json(AgentStatusResponse {
                agent_id,
                status: format!("{:?}", status),
            }))
        }
        None => Err((
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: format!("Agent '{}' not found", agent_id),
            }),
        )),
    }
}

// --- Token endpoints ---

#[derive(Serialize)]
pub struct TokenReport {
    pub message: String,
}

pub async fn token_report() -> Json<TokenReport> {
    Json(TokenReport {
        message: "Token reporting — detailed per-agent usage coming soon".to_string(),
    })
}

// --- WebSocket streaming endpoint ---

#[derive(Deserialize)]
struct WsCommand {
    agent_id: String,
    input: String,
}

pub async fn ws_handler(
    ws: axum::extract::WebSocketUpgrade,
    State(state): State<AppState>,
) -> axum::response::Response {
    ws.on_upgrade(move |socket| handle_ws(socket, state))
}

async fn handle_ws(mut socket: axum::extract::ws::WebSocket, state: AppState) {
    use axum::extract::ws::Message;
    use axocoatl_llm::StreamEvent;
    use tokio_stream::StreamExt;

    while let Some(Ok(msg)) = socket.recv().await {
        match msg {
            Message::Text(text) => {
                let cmd: Result<WsCommand, _> = serde_json::from_str(&text);
                match cmd {
                    Ok(cmd) => {
                        let daemon = state.read().await;

                        // Try streaming first via the provider directly
                        let provider = daemon
                            .provider_registry
                            .get("openai")
                            .or_else(|| daemon.provider_registry.get("anthropic"));

                        match provider {
                            Some(provider) => {
                                let request = axocoatl_llm::ChatRequest::simple(&cmd.input);
                                match provider.chat_stream(request).await {
                                    Ok(mut stream) => {
                                        drop(daemon); // Release read lock during streaming
                                        while let Some(event) = stream.next().await {
                                            let ws_msg = match event {
                                                Ok(StreamEvent::TextDelta { delta }) => {
                                                    serde_json::json!({
                                                        "type": "text_delta",
                                                        "delta": delta
                                                    })
                                                }
                                                Ok(StreamEvent::Done { finish_reason }) => {
                                                    serde_json::json!({
                                                        "type": "done",
                                                        "finish_reason": format!("{:?}", finish_reason)
                                                    })
                                                }
                                                Ok(StreamEvent::Usage(usage)) => {
                                                    serde_json::json!({
                                                        "type": "usage",
                                                        "input_tokens": usage.input_tokens,
                                                        "output_tokens": usage.output_tokens
                                                    })
                                                }
                                                Ok(_) => continue,
                                                Err(e) => {
                                                    serde_json::json!({
                                                        "type": "error",
                                                        "error": e.to_string()
                                                    })
                                                }
                                            };

                                            if socket
                                                .send(Message::Text(ws_msg.to_string().into()))
                                                .await
                                                .is_err()
                                            {
                                                break;
                                            }
                                        }
                                    }
                                    Err(_) => {
                                        // Streaming not supported — fall back to non-streaming
                                        match daemon.execute_agent(&cmd.agent_id, &cmd.input).await
                                        {
                                            Ok(output) => {
                                                let msg = serde_json::json!({
                                                    "type": "text_delta",
                                                    "delta": output
                                                });
                                                let _ = socket
                                                    .send(Message::Text(msg.to_string().into()))
                                                    .await;
                                                let done = serde_json::json!({"type": "done", "finish_reason": "Stop"});
                                                let _ = socket
                                                    .send(Message::Text(done.to_string().into()))
                                                    .await;
                                            }
                                            Err(e) => {
                                                let err = serde_json::json!({"type": "error", "error": e.to_string()});
                                                let _ = socket
                                                    .send(Message::Text(err.to_string().into()))
                                                    .await;
                                            }
                                        }
                                    }
                                }
                            }
                            None => {
                                let err = serde_json::json!({"type": "error", "error": "No provider available"});
                                let _ = socket.send(Message::Text(err.to_string().into())).await;
                            }
                        }
                    }
                    Err(e) => {
                        let err = serde_json::json!({"type": "error", "error": format!("Invalid command: {e}")});
                        let _ = socket.send(Message::Text(err.to_string().into())).await;
                    }
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }
}