helios-remote/crates/server/src/api.rs

use std::time::Duration;
use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use tracing::error;

use helios_common::protocol::{ClientMessage, ServerMessage};
use crate::AppState;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

// ── Response types ──────────────────────────────────────────────────────────

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

fn not_found(label: &str) -> (StatusCode, Json<ErrorBody>) {
    (
        StatusCode::NOT_FOUND,
        Json(ErrorBody {
            error: format!("Device '{label}' not found or not connected"),
        }),
    )
}

fn timeout_error(label: &str, op: &str) -> (StatusCode, Json<ErrorBody>) {
    (
        StatusCode::GATEWAY_TIMEOUT,
        Json(ErrorBody {
            error: format!("Timed out waiting for client response (device='{label}', op='{op}')"),
        }),
    )
}

fn send_error(label: &str, op: &str) -> (StatusCode, Json<ErrorBody>) {
    (
        StatusCode::BAD_GATEWAY,
        Json(ErrorBody {
            error: format!("Failed to send command to client — may have disconnected (device='{label}', op='{op}')"),
        }),
    )
}

// ── Helper to send a command and await the response ─────────────────────────

async fn dispatch<F>(
    state: &AppState,
    label: &str,
    op: &str,
    make_msg: F,
) -> Result<ClientMessage, (StatusCode, Json<ErrorBody>)>
where
    F: FnOnce(Uuid) -> ServerMessage,
{
    dispatch_with_timeout(state, label, op, make_msg, REQUEST_TIMEOUT).await
}

async fn dispatch_with_timeout<F>(
    state: &AppState,
    label: &str,
    op: &str,
    make_msg: F,
    timeout: Duration,
) -> Result<ClientMessage, (StatusCode, Json<ErrorBody>)>
where
    F: FnOnce(Uuid) -> ServerMessage,
{
    let tx = state
        .sessions
        .get_cmd_tx(label)
        .ok_or_else(|| not_found(label))?;

    let request_id = Uuid::new_v4();
    let rx = state.sessions.register_pending(request_id);
    let msg = make_msg(request_id);

    tx.send(msg).await.map_err(|e| {
        error!("Channel send failed for device={label}, op={op}: {e}");
        send_error(label, op)
    })?;

    match tokio::time::timeout(timeout, rx).await {
        Ok(Ok(response)) => Ok(response),
        Ok(Err(_)) => Err(send_error(label, op)),
        Err(_) => Err(timeout_error(label, op)),
    }
}

// ── Handlers ─────────────────────────────────────────────────────────────────

/// GET /devices — list all connected clients
pub async fn list_devices(State(state): State<AppState>) -> Json<serde_json::Value> {
    let devices = state.sessions.list();
    Json(serde_json::json!({ "devices": devices }))
}

/// POST /devices/:label/screenshot — full screen screenshot
pub async fn request_screenshot(
    Path(label): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "screenshot", |rid| {
        ServerMessage::ScreenshotRequest { request_id: rid }
    }).await {
        Ok(ClientMessage::ScreenshotResponse { image_base64, width, height, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "image_base64": image_base64, "width": width, "height": height })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/windows/:window_id/screenshot
pub async fn window_screenshot(
    Path((label, window_id)): Path<(String, u64)>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "window_screenshot", |rid| {
        ServerMessage::WindowScreenshotRequest { request_id: rid, window_id }
    }).await {
        Ok(ClientMessage::ScreenshotResponse { image_base64, width, height, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "image_base64": image_base64, "width": width, "height": height })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// GET /devices/:label/logs?lines=100
pub async fn logs(
    Path(label): Path<String>,
    Query(params): Query<std::collections::HashMap<String, String>>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    let lines: u32 = params.get("lines").and_then(|v| v.parse().ok()).unwrap_or(100);
    match dispatch(&state, &label, "logs", |rid| {
        ServerMessage::LogsRequest { request_id: rid, lines }
    }).await {
        Ok(ClientMessage::LogsResponse { content, log_path, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "content": content, "log_path": log_path, "lines": lines })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/exec
#[derive(Deserialize)]
pub struct ExecBody {
    pub command: String,
    pub timeout_ms: Option<u64>,
}

pub async fn request_exec(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Json(body): Json<ExecBody>,
) -> impl IntoResponse {
    let server_timeout = body.timeout_ms
        .map(|ms| Duration::from_millis(ms + 5_000))
        .unwrap_or(REQUEST_TIMEOUT);

    match dispatch_with_timeout(&state, &label, "exec", |rid| ServerMessage::ExecRequest {
        request_id: rid,
        command: body.command.clone(),
        timeout_ms: body.timeout_ms,
    }, server_timeout).await {
        Ok(ClientMessage::ExecResponse { stdout, stderr, exit_code, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "stdout": stdout, "stderr": stderr, "exit_code": exit_code })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// GET /devices/:label/windows
pub async fn list_windows(
    Path(label): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "list_windows", |rid| {
        ServerMessage::ListWindowsRequest { request_id: rid }
    }).await {
        Ok(ClientMessage::ListWindowsResponse { windows, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "windows": windows })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/windows/minimize-all
pub async fn minimize_all(
    Path(label): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "minimize_all", |rid| {
        ServerMessage::MinimizeAllRequest { request_id: rid }
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/windows/:window_id/focus
pub async fn focus_window(
    Path((label, window_id)): Path<(String, u64)>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "focus_window", |rid| {
        ServerMessage::FocusWindowRequest { request_id: rid, window_id }
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/windows/:window_id/maximize
pub async fn maximize_and_focus(
    Path((label, window_id)): Path<(String, u64)>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "maximize_and_focus", |rid| {
        ServerMessage::MaximizeAndFocusRequest { request_id: rid, window_id }
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// GET /version — server version (public, no auth)
pub async fn server_version() -> impl IntoResponse {
    Json(serde_json::json!({
        "commit": env!("GIT_COMMIT"),
    }))
}

/// GET /devices/:label/version — client version
pub async fn client_version(
    Path(label): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "version", |rid| {
        ServerMessage::VersionRequest { request_id: rid }
    }).await {
        Ok(ClientMessage::VersionResponse { version, commit, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "version": version, "commit": commit })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/upload
#[derive(Deserialize)]
pub struct UploadBody {
    pub path: String,
    pub content_base64: String,
}

pub async fn upload_file(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Json(body): Json<UploadBody>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "upload", |rid| ServerMessage::UploadRequest {
        request_id: rid,
        path: body.path.clone(),
        content_base64: body.content_base64.clone(),
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// GET /devices/:label/download?path=...
#[derive(Deserialize)]
pub struct DownloadQuery {
    pub path: String,
}

pub async fn download_file(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Query(query): Query<DownloadQuery>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "download", |rid| ServerMessage::DownloadRequest {
        request_id: rid,
        path: query.path.clone(),
    }).await {
        Ok(ClientMessage::DownloadResponse { content_base64, size, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "content_base64": content_base64, "size": size })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/run
#[derive(Deserialize)]
pub struct RunBody {
    pub program: String,
    #[serde(default)]
    pub args: Vec<String>,
}

pub async fn run_program(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Json(body): Json<RunBody>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "run", |rid| ServerMessage::RunRequest {
        request_id: rid,
        program: body.program.clone(),
        args: body.args.clone(),
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// GET /devices/:label/clipboard
pub async fn clipboard_get(
    Path(label): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "clipboard_get", |rid| {
        ServerMessage::ClipboardGetRequest { request_id: rid }
    }).await {
        Ok(ClientMessage::ClipboardGetResponse { text, .. }) => (
            StatusCode::OK,
            Json(serde_json::json!({ "text": text })),
        ).into_response(),
        Ok(ClientMessage::Error { message, .. }) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": message })),
        ).into_response(),
        Ok(_) => (StatusCode::BAD_GATEWAY, Json(serde_json::json!({ "error": "Unexpected response" }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/clipboard
#[derive(Deserialize)]
pub struct ClipboardSetBody {
    pub text: String,
}

pub async fn clipboard_set(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Json(body): Json<ClipboardSetBody>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "clipboard_set", |rid| {
        ServerMessage::ClipboardSetRequest { request_id: rid, text: body.text.clone() }
    }).await {
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}

/// POST /devices/:label/prompt
#[derive(Deserialize)]
pub struct PromptBody {
    pub message: String,
    pub title: Option<String>,
}

pub async fn prompt_user(
    Path(label): Path<String>,
    State(state): State<AppState>,
    Json(body): Json<PromptBody>,
) -> impl IntoResponse {
    match dispatch(&state, &label, "prompt", |rid| ServerMessage::PromptRequest {
        request_id: rid,
        message: body.message.clone(),
        title: body.title.clone(),
    }).await {
        Ok(ClientMessage::PromptResponse { answer, .. }) => {
            (StatusCode::OK, Json(serde_json::json!({ "ok": true, "answer": answer }))).into_response()
        }
        Ok(_) => (StatusCode::OK, Json(serde_json::json!({ "ok": true }))).into_response(),
        Err(e) => e.into_response(),
    }
}