bot/src/domain/tasks.rs

use chrono::Utc;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, 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::entities::{push_subscription, task_subscription};
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.clone());

    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(),
        ));

    // Send Push Notifications to subscribers
    let subscriptions = task_subscription::Entity::find()
        .filter(task_subscription::Column::TaskId.eq(task_id))
        .all(db)
        .await
        .map_err(crate::error::AppError::Database)?;

    tracing::info!(
        "Found {} task subscriptions for task {}",
        subscriptions.len(),
        task_id
    );

    for sub in subscriptions {
        let push_subs = push_subscription::Entity::find()
            .filter(push_subscription::Column::UserSub.eq(sub.user_sub.clone()))
            .all(db)
            .await
            .map_err(crate::error::AppError::Database)?;

        tracing::info!(
            "Found {} push subscriptions for user {}",
            push_subs.len(),
            sub.user_sub
        );

        for push_sub in push_subs {
            let sender = _scheduler.push_sender.clone();
            let goal = task_response.goal.clone();
            let status = status.clone();
            let sub_data = crate::domain::notifications::push::PushSubscription {
                endpoint: push_sub.endpoint,
                p256dh: push_sub.p256dh,
                auth: push_sub.auth,
            };

            tokio::spawn(async move {
                if let Err(e) = sender
                    .send_notification(&sub_data, &format!("Task Completed: {}", status), &goal)
                    .await
                {
                    tracing::error!("Failed to send notification in background task: {}", e);
                }
            });
        }
    }

    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(),
    })
}