lightOps/crates/lightops-server/src/task.rs

use crate::{
    error::AppError,
    state::{AppState, TaskReply},
};
use chrono::Utc;
use lightops_common::protocol::ServerMessage;
use serde_json::Value;
use tokio::sync::oneshot;
use uuid::Uuid;

pub async fn request_agent_task(
    state: &AppState,
    agent_id: &str,
    user_id: Option<i64>,
    action: &str,
    params: Value,
) -> Result<TaskReply, AppError> {
    request_agent_task_with_timeout(
        state,
        agent_id,
        user_id,
        action,
        params,
        state.cfg.task_timeout_secs,
    )
    .await
}

pub async fn request_agent_task_with_timeout(
    state: &AppState,
    agent_id: &str,
    user_id: Option<i64>,
    action: &str,
    params: Value,
    timeout_secs: u64,
) -> Result<TaskReply, AppError> {
    request_agent_task_with_timeout_and_stored_params(
        state,
        agent_id,
        user_id,
        action,
        params.clone(),
        params,
        timeout_secs,
    )
    .await
}

pub async fn request_agent_task_with_timeout_and_stored_params(
    state: &AppState,
    agent_id: &str,
    user_id: Option<i64>,
    action: &str,
    params: Value,
    stored_params: Value,
    timeout_secs: u64,
) -> Result<TaskReply, AppError> {
    let task_id = Uuid::new_v4().to_string();
    sqlx::query(
        "INSERT INTO tasks (id, agent_id, user_id, action, params_json, status, started_at) VALUES (?, ?, ?, ?, ?, 'running', ?)",
    )
    .bind(&task_id)
    .bind(agent_id)
    .bind(user_id)
    .bind(action)
    .bind(stored_params.to_string())
    .bind(Utc::now().to_rfc3339())
    .execute(&state.pool)
    .await?;
    append_task_event(
        state,
        &task_id,
        "info",
        "任务已创建",
        Some(&serde_json::json!({ "action": action })),
    )
    .await?;

    let handle = {
        let agents = state.agents.read().await;
        agents.get(agent_id).cloned()
    }
    .ok_or(AppError::AgentOffline)?;

    let (tx, rx) = oneshot::channel();
    state.pending.insert(task_id.clone(), tx);
    state
        .pending_agents
        .insert(task_id.clone(), agent_id.to_string());
    let msg = ServerMessage::TaskRequest {
        task_id: task_id.clone(),
        action: action.to_string(),
        params,
    };
    if handle.tx.send(msg).is_err() {
        state.pending.remove(&task_id);
        state.pending_agents.remove(&task_id);
        mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?;
        return Err(AppError::AgentOffline);
    }
    append_task_event(state, &task_id, "info", "任务已下发到 Agent", None).await?;

    let reply = match tokio::time::timeout(std::time::Duration::from_secs(timeout_secs), rx).await {
        Ok(Ok(reply)) => reply,
        Ok(Err(_)) => {
            state.pending_agents.remove(&task_id);
            mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?;
            return Err(AppError::AgentOffline);
        }
        Err(_) => {
            state.pending.remove(&task_id);
            state.pending_agents.remove(&task_id);
            mark_task(state, &task_id, "timeout", None, Some("任务超时")).await?;
            return Err(AppError::Timeout);
        }
    };

    let status = if reply.success { "success" } else { "failed" };
    mark_task(
        state,
        &task_id,
        status,
        Some(&reply.data),
        reply.error.as_deref(),
    )
    .await?;
    Ok(reply)
}

pub async fn append_task_event(
    state: &AppState,
    task_id: &str,
    level: &str,
    message: &str,
    data: Option<&Value>,
) -> Result<(), AppError> {
    sqlx::query(
        "INSERT INTO task_events(task_id, level, message, data_json, created_at) VALUES(?, ?, ?, ?, ?)",
    )
    .bind(task_id)
    .bind(level)
    .bind(message.chars().take(500).collect::<String>())
    .bind(data.map(|value| value.to_string()))
    .bind(Utc::now().to_rfc3339())
    .execute(&state.pool)
    .await?;
    Ok(())
}

pub async fn fail_agent_pending_tasks(state: &AppState, agent_id: &str, reason: &str) {
    let task_ids: Vec<String> = state
        .pending_agents
        .iter()
        .filter(|item| item.value() == agent_id)
        .map(|item| item.key().clone())
        .collect();

    for task_id in task_ids {
        state.pending_agents.remove(&task_id);
        if let Some((_, tx)) = state.pending.remove(&task_id) {
            let _ = tx.send(TaskReply {
                success: false,
                data: serde_json::json!({}),
                error: Some(reason.to_string()),
            });
        }
        let _ = mark_task(state, &task_id, "failed", None, Some(reason)).await;
    }
}

pub async fn mark_task(
    state: &AppState,
    task_id: &str,
    status: &str,
    result: Option<&Value>,
    error: Option<&str>,
) -> Result<(), AppError> {
    sqlx::query(
        "UPDATE tasks SET status = ?, result_json = ?, error = ?, finished_at = ? WHERE id = ?",
    )
    .bind(status)
    .bind(result.map(|v| v.to_string()))
    .bind(error)
    .bind(Utc::now().to_rfc3339())
    .bind(task_id)
    .execute(&state.pool)
    .await?;
    let level = if status == "success" { "info" } else { "error" };
    let message = match status {
        "success" => "任务成功完成",
        "timeout" => "任务执行超时",
        "cancelled" => "任务已取消",
        _ => "任务执行失败",
    };
    let data = serde_json::json!({
        "status": status,
        "error": error,
        "result": result
    });
    append_task_event(state, task_id, level, message, Some(&data)).await?;
    Ok(())
}

pub async fn audit(
    state: &AppState,
    user_id: Option<i64>,
    agent_id: Option<&str>,
    action: &str,
    target: Option<&str>,
    params_summary: Option<String>,
    success: bool,
    error: Option<&str>,
) {
    let _ = sqlx::query(
        "INSERT INTO audit_logs (user_id, agent_id, action, target, params_summary, success, error) VALUES (?, ?, ?, ?, ?, ?, ?)",
    )
    .bind(user_id)
    .bind(agent_id)
    .bind(action)
    .bind(target)
    .bind(params_summary)
    .bind(if success { 1 } else { 0 })
    .bind(error)
    .execute(&state.pool)
    .await;
}