bot/src/domain/tasks.rs

use chrono::Utc;
use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, Set};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use uuid::Uuid;

use crate::config::Config;
use crate::domain::agent::Agent;
use crate::entities::task::Entity as Task;
use crate::entities::task_run::{self, Entity as TaskRun};
use crate::scheduler::Scheduler;

#[derive(Deserialize)]
pub struct CreateTaskRequest {
    pub goal: String,
    pub cron: Option<String>,
}

#[derive(Deserialize)]
pub struct UpdateTaskRequest {
    pub goal: String,
    pub cron: Option<String>,
}

#[derive(Serialize, Clone, Debug)]
pub struct TaskResponse {
    pub id: Uuid,
    pub goal: String,
    pub cron: Option<String>,
    pub created_at: chrono::DateTime<chrono::FixedOffset>,
    pub runs: Vec<TaskRunResponse>,
}

#[derive(Serialize, Clone, Debug)]
pub struct TaskRunResponse {
    pub id: Uuid,
    pub status: String,
    pub logs: String,
    pub answer: Option<String>,
    pub created_at: chrono::DateTime<chrono::FixedOffset>,
}

#[derive(Serialize, Clone, Debug)]
pub struct RecentRunResponse {
    pub id: Uuid,
    pub task_id: Uuid,
    pub goal: String,
    pub status: String,
    pub created_at: chrono::DateTime<chrono::FixedOffset>,
}

use crate::error::AppResult;

pub async fn execute_agent_run(
    db: &DatabaseConnection,
    _scheduler: &Arc<Scheduler>,
    config: &Arc<Config>,
    task_id: Uuid,
    goal: String,
) -> AppResult<TaskResponse> {
    let run_id = Uuid::new_v4();
    tracing::info!(%task_id, %run_id, "Starting agent execution run");

    let new_run = task_run::ActiveModel {
        id: Set(run_id),
        task_id: Set(task_id),
        status: Set("running".to_string()),
        logs: Set(String::new()),
        answer: Set(None),
        created_at: Set(Utc::now().into()),
    };

    new_run
        .insert(db)
        .await
        .map_err(crate::error::AppError::Database)?;

    let _ = _scheduler
        .tx
        .send(crate::server::notifications::WsEvent::RunStarted(
            RecentRunResponse {
                id: run_id,
                task_id,
                goal: goal.clone(),
                status: "running".to_string(),
                created_at: Utc::now().into(),
            },
        ));

    let mut agent = Agent::new(
        db.clone(),
        config.zen_api_key.clone(),
        config.tavily_api_key.clone(),
        goal.clone(),
    )?;

    let (logs, answer, status) = match agent.run(config).await {
        Ok((logs, answer)) => {
            tracing::info!(%task_id, %run_id, "Agent execution completed successfully");
            (logs, answer, "completed".to_string())
        }
        Err(e) => {
            tracing::error!(%task_id, %run_id, error = %e, "Agent execution failed");
            (
                format!("Execution failed: {}", e),
                None,
                "failed".to_string(),
            )
        }
    };

    let run: task_run::ActiveModel = TaskRun::find_by_id(run_id)
        .one(db)
        .await
        .map_err(crate::error::AppError::Database)?
        .ok_or_else(|| crate::error::AppError::NotFound("Run not found after insert".into()))?
        .into();

    let mut run = run;
    run.logs = Set(logs.clone());
    run.answer = Set(answer.clone());
    run.status = Set(status);

    run.update(db)
        .await
        .map_err(crate::error::AppError::Database)?;

    let task_response = get_task_inner(task_id, db).await?;
    let _ = _scheduler
        .tx
        .send(crate::server::notifications::WsEvent::RunFinished(
            task_response.clone(),
        ));

    Ok(task_response)
}

pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult<TaskResponse> {
    let results = Task::find_by_id(id)
        .find_with_related(TaskRun)
        .all(db)
        .await
        .map_err(crate::error::AppError::Database)?;

    let (t, mut runs) = results
        .into_iter()
        .next()
        .ok_or_else(|| crate::error::AppError::NotFound("Task not found".into()))?;

    runs.sort_by(|a, b| b.created_at.cmp(&a.created_at));

    Ok(TaskResponse {
        id: t.id,
        goal: t.goal,
        cron: t.cron,
        created_at: t.created_at,
        runs: runs
            .into_iter()
            .map(|r| TaskRunResponse {
                id: r.id,
                status: r.status,
                logs: r.logs,
                answer: r.answer,
                created_at: r.created_at,
            })
            .collect(),
    })
}