PicoBot/src/tools/scheduler_manage.rs

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::json;

use crate::config::SchedulerSchedule;
use crate::storage::{
    SchedulerJobRecord, SchedulerJobState, SchedulerJobUpsert, SessionStore,
};
use crate::tools::traits::{Tool, ToolResult};

pub struct SchedulerManageTool {
    store: Arc<SessionStore>,
}

impl SchedulerManageTool {
    pub fn new(store: Arc<SessionStore>) -> Self {
        Self { store }
    }
}

#[async_trait]
impl Tool for SchedulerManageTool {
    fn name(&self) -> &str {
        "scheduler_manage"
    }

    fn description(&self) -> &str {
        "Manage DB-backed scheduled jobs. Supports actions: list, get, put, delete, pause, resume. Jobs persist in SQLite and are executed by the scheduler runtime."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["list", "get", "put", "delete", "pause", "resume"]
                },
                "id": {
                    "type": "string",
                    "description": "Job id"
                },
                "enabled_only": {
                    "type": "boolean",
                    "description": "Only for list action"
                },
                "kind": {
                    "type": "string",
                    "enum": ["internal_event", "outbound_message", "agent_task"]
                },
                "schedule": {
                    "type": "object",
                    "description": "Schedule object, for example {type: 'interval', seconds: 300} or {type: 'cron', expression: '0 9 * * *'}"
                },
                "target": {
                    "type": "object"
                },
                "payload": {
                    "type": "object",
                    "description": "Job payload. agent_task supports prompt, fresh_session, system_prompt, sender_id, metadata. outbound_message expects content. internal_event expects event."
                },
                "max_runs": {
                    "type": ["integer", "null"]
                }
            },
            "required": ["action"]
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let action = match args.get("action").and_then(|value| value.as_str()) {
            Some(action) => action,
            None => return Ok(error_result("Missing required parameter: action")),
        };

        let output = match action {
            "list" => {
                let enabled_only = args
                    .get("enabled_only")
                    .and_then(|value| value.as_bool())
                    .unwrap_or(false);
                let jobs = self.store.list_scheduler_jobs(enabled_only)?;
                json!(jobs.iter().map(record_to_json).collect::<Vec<_>>())
            }
            "get" => {
                let id = require_str(&args, "id")?;
                match self.store.get_scheduler_job(id)? {
                    Some(record) => record_to_json(&record),
                    None => return Ok(error_result(&format!("scheduler job '{}' not found", id))),
                }
            }
            "put" => {
                let input = build_upsert(&args)?;
                let record = self.store.upsert_scheduler_job(&input)?;
                record_to_json(&record)
            }
            "delete" => {
                let id = require_str(&args, "id")?;
                self.store.delete_scheduler_job(id)?;
                json!({"status": "deleted", "id": id})
            }
            "pause" => {
                let id = require_str(&args, "id")?;
                let record = self
                    .store
                    .get_scheduler_job(id)?
                    .ok_or_else(|| anyhow::anyhow!("scheduler job '{}' not found", id))?;
                let mut input = record_to_upsert(&record);
                input.enabled = false;
                input.state = SchedulerJobState::Paused;
                input.paused_at = Some(current_timestamp());
                input.next_fire_at = None;
                let saved = self.store.upsert_scheduler_job(&input)?;
                record_to_json(&saved)
            }
            "resume" => {
                let id = require_str(&args, "id")?;
                let record = self
                    .store
                    .get_scheduler_job(id)?
                    .ok_or_else(|| anyhow::anyhow!("scheduler job '{}' not found", id))?;
                let mut input = record_to_upsert(&record);
                input.enabled = true;
                input.state = SchedulerJobState::Scheduled;
                input.paused_at = None;
                input.completed_at = None;
                input.next_fire_at = None;
                let saved = self.store.upsert_scheduler_job(&input)?;
                record_to_json(&saved)
            }
            _ => return Ok(error_result("Unsupported action")),
        };

        Ok(ToolResult {
            success: true,
            output: serde_json::to_string_pretty(&output)?,
            error: None,
        })
    }
}

fn build_upsert(args: &serde_json::Value) -> anyhow::Result<SchedulerJobUpsert> {
    let id = require_str(args, "id")?.to_string();
    let kind = require_str(args, "kind")?.to_string();
    let schedule_value = args
        .get("schedule")
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("Missing required parameter: schedule"))?;
    let schedule: SchedulerSchedule = serde_json::from_value(schedule_value.clone())?;
    schedule.validate(&id)?;

    let (interval_secs, startup_delay_secs) = match &schedule {
        SchedulerSchedule::Interval {
            seconds,
            startup_delay_secs,
        } => (*seconds as i64, *startup_delay_secs as i64),
        _ => (0, 0),
    };

    Ok(SchedulerJobUpsert {
        id,
        kind,
        schedule: serde_json::to_value(schedule)?,
        interval_secs,
        startup_delay_secs,
        target: args.get("target").cloned().unwrap_or_else(|| json!({})),
        payload: args.get("payload").cloned().unwrap_or_else(|| json!({})),
        enabled: args.get("enabled").and_then(|value| value.as_bool()).unwrap_or(true),
        state: if args.get("enabled").and_then(|value| value.as_bool()).unwrap_or(true) {
            SchedulerJobState::Scheduled
        } else {
            SchedulerJobState::Paused
        },
        last_status: None,
        last_error: None,
        run_count: 0,
        max_runs: args.get("max_runs").and_then(|value| value.as_i64()),
        last_fired_at: None,
        next_fire_at: None,
        paused_at: None,
        completed_at: None,
    })
}

fn record_to_json(record: &SchedulerJobRecord) -> serde_json::Value {
    json!({
        "id": record.id,
        "kind": record.kind,
        "schedule": record.schedule,
        "target": record.target,
        "payload": record.payload,
        "enabled": record.enabled,
        "state": record.state,
        "last_status": record.last_status,
        "last_error": record.last_error,
        "run_count": record.run_count,
        "max_runs": record.max_runs,
        "last_fired_at": record.last_fired_at,
        "next_fire_at": record.next_fire_at,
        "paused_at": record.paused_at,
        "completed_at": record.completed_at,
        "created_at": record.created_at,
        "updated_at": record.updated_at,
    })
}

fn record_to_upsert(record: &SchedulerJobRecord) -> SchedulerJobUpsert {
    SchedulerJobUpsert {
        id: record.id.clone(),
        kind: record.kind.clone(),
        schedule: record.schedule.clone(),
        interval_secs: record.interval_secs,
        startup_delay_secs: record.startup_delay_secs,
        target: record.target.clone(),
        payload: record.payload.clone(),
        enabled: record.enabled,
        state: record.state.clone(),
        last_status: record.last_status.clone(),
        last_error: record.last_error.clone(),
        run_count: record.run_count,
        max_runs: record.max_runs,
        last_fired_at: record.last_fired_at,
        next_fire_at: record.next_fire_at,
        paused_at: record.paused_at,
        completed_at: record.completed_at,
    }
}

fn require_str<'a>(args: &'a serde_json::Value, key: &str) -> anyhow::Result<&'a str> {
    args.get(key)
        .and_then(|value| value.as_str())
        .ok_or_else(|| anyhow::anyhow!("Missing required parameter: {}", key))
}

fn error_result(message: &str) -> ToolResult {
    ToolResult {
        success: false,
        output: String::new(),
        error: Some(message.to_string()),
    }
}

fn current_timestamp() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as i64
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_scheduler_manage_put_and_get() {
        let store = Arc::new(SessionStore::in_memory().unwrap());
        let tool = SchedulerManageTool::new(store);

        let put_result = tool
            .execute(json!({
                "action": "put",
                "id": "heartbeat",
                "kind": "outbound_message",
                "schedule": {
                    "type": "interval",
                    "seconds": 60
                },
                "target": {
                    "channel": "feishu",
                    "chat_id": "oc_demo"
                },
                "payload": {
                    "content": "ping"
                }
            }))
            .await
            .unwrap();
        assert!(put_result.success);

        let get_result = tool
            .execute(json!({
                "action": "get",
                "id": "heartbeat"
            }))
            .await
            .unwrap();
        assert!(get_result.success);
        assert!(get_result.output.contains("heartbeat"));
        assert!(get_result.output.contains("outbound_message"));
    }

    #[tokio::test]
    async fn test_scheduler_manage_put_agent_task() {
        let store = Arc::new(SessionStore::in_memory().unwrap());
        let tool = SchedulerManageTool::new(store);

        let put_result = tool
            .execute(json!({
                "action": "put",
                "id": "agent.daily_summary",
                "kind": "agent_task",
                "schedule": {
                    "type": "cron",
                    "expression": "0 9 * * *"
                },
                "target": {
                    "channel": "feishu",
                    "chat_id": "oc_demo"
                },
                "payload": {
                    "prompt": "请总结今天待办"
                }
            }))
            .await
            .unwrap();

        assert!(put_result.success);
        assert!(put_result.output.contains("agent_task"));
    }
}