use std::collections::HashMap; use serde::{Deserialize, Serialize}; use anyhow::{Result, Context}; use chrono::{DateTime, Utc, Duration}; use crate::config::Config; use crate::persona::Persona; use crate::transmission::TransmissionController; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ScheduledTask { pub id: String, pub task_type: TaskType, pub next_run: DateTime, pub interval_hours: Option, pub enabled: bool, pub last_run: Option>, pub run_count: u32, pub max_runs: Option, pub created_at: DateTime, pub metadata: HashMap, } #[derive(Debug, Clone, Serialize, Deserialize)] pub enum TaskType { DailyMaintenance, AutoTransmission, RelationshipDecay, BreakthroughCheck, MaintenanceTransmission, Custom(String), } impl std::fmt::Display for TaskType { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { TaskType::DailyMaintenance => write!(f, "daily_maintenance"), TaskType::AutoTransmission => write!(f, "auto_transmission"), TaskType::RelationshipDecay => write!(f, "relationship_decay"), TaskType::BreakthroughCheck => write!(f, "breakthrough_check"), TaskType::MaintenanceTransmission => write!(f, "maintenance_transmission"), TaskType::Custom(name) => write!(f, "custom_{}", name), } } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TaskExecution { pub task_id: String, pub execution_time: DateTime, pub duration_ms: u64, pub success: bool, pub result: Option, pub error: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AIScheduler { config: Config, tasks: HashMap, execution_history: Vec, last_check: Option>, } impl AIScheduler { pub fn new(config: &Config) -> Result { let (tasks, execution_history) = Self::load_scheduler_data(config)?; let mut scheduler = AIScheduler { config: config.clone(), tasks, execution_history, last_check: None, }; // Initialize default tasks if none exist if scheduler.tasks.is_empty() { scheduler.create_default_tasks()?; } Ok(scheduler) } pub async fn run_scheduled_tasks(&mut self, persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result> { let now = Utc::now(); let mut executions = Vec::new(); // Find tasks that are due to run let due_task_ids: Vec = self.tasks .iter() .filter(|(_, task)| task.enabled && task.next_run <= now) .filter(|(_, task)| { // Check if task hasn't exceeded max runs if let Some(max_runs) = task.max_runs { task.run_count < max_runs } else { true } }) .map(|(id, _)| id.clone()) .collect(); for task_id in due_task_ids { let execution = self.execute_task(&task_id, persona, transmission_controller).await?; executions.push(execution); } self.last_check = Some(now); self.save_scheduler_data()?; Ok(executions) } async fn execute_task(&mut self, task_id: &str, persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result { let start_time = Utc::now(); let mut execution = TaskExecution { task_id: task_id.to_string(), execution_time: start_time, duration_ms: 0, success: false, result: None, error: None, }; // Get task type without borrowing mutably let task_type = { let task = self.tasks.get(task_id) .ok_or_else(|| anyhow::anyhow!("Task not found: {}", task_id))?; task.task_type.clone() }; // Execute the task based on its type let result = match &task_type { TaskType::DailyMaintenance => self.execute_daily_maintenance(persona, transmission_controller).await, TaskType::AutoTransmission => self.execute_auto_transmission(persona, transmission_controller).await, TaskType::RelationshipDecay => self.execute_relationship_decay(persona).await, TaskType::BreakthroughCheck => self.execute_breakthrough_check(persona, transmission_controller).await, TaskType::MaintenanceTransmission => self.execute_maintenance_transmission(persona, transmission_controller).await, TaskType::Custom(name) => self.execute_custom_task(name, persona, transmission_controller).await, }; let end_time = Utc::now(); execution.duration_ms = (end_time - start_time).num_milliseconds() as u64; // Now update the task state with mutable borrow match result { Ok(message) => { execution.success = true; execution.result = Some(message); // Update task state if let Some(task) = self.tasks.get_mut(task_id) { task.last_run = Some(start_time); task.run_count += 1; // Schedule next run if recurring if let Some(interval_hours) = task.interval_hours { task.next_run = start_time + Duration::hours(interval_hours); } else { // One-time task, disable it task.enabled = false; } } } Err(e) => { execution.error = Some(e.to_string()); // For failed tasks, retry in a shorter interval if let Some(task) = self.tasks.get_mut(task_id) { if task.interval_hours.is_some() { task.next_run = start_time + Duration::minutes(15); // Retry in 15 minutes } } } } self.execution_history.push(execution.clone()); // Keep only recent execution history (last 1000 executions) if self.execution_history.len() > 1000 { self.execution_history.drain(..self.execution_history.len() - 1000); } Ok(execution) } async fn execute_daily_maintenance(&self, persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result { // Run daily maintenance persona.daily_maintenance()?; // Check for maintenance transmissions let transmissions = transmission_controller.check_maintenance_transmissions(persona).await?; Ok(format!("Daily maintenance completed. {} maintenance transmissions sent.", transmissions.len())) } async fn execute_auto_transmission(&self, _persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result { let transmissions = transmission_controller.check_autonomous_transmissions(_persona).await?; Ok(format!("Autonomous transmission check completed. {} transmissions sent.", transmissions.len())) } async fn execute_relationship_decay(&self, persona: &mut Persona) -> Result { persona.daily_maintenance()?; Ok("Relationship time decay applied.".to_string()) } async fn execute_breakthrough_check(&self, persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result { let transmissions = transmission_controller.check_breakthrough_transmissions(persona).await?; Ok(format!("Breakthrough check completed. {} transmissions sent.", transmissions.len())) } async fn execute_maintenance_transmission(&self, persona: &mut Persona, transmission_controller: &mut TransmissionController) -> Result { let transmissions = transmission_controller.check_maintenance_transmissions(persona).await?; Ok(format!("Maintenance transmission check completed. {} transmissions sent.", transmissions.len())) } async fn execute_custom_task(&self, _name: &str, _persona: &mut Persona, _transmission_controller: &mut TransmissionController) -> Result { // Placeholder for custom task execution Ok("Custom task executed.".to_string()) } pub fn create_task(&mut self, task_type: TaskType, next_run: DateTime, interval_hours: Option) -> Result { let task_id = uuid::Uuid::new_v4().to_string(); let now = Utc::now(); let task = ScheduledTask { id: task_id.clone(), task_type, next_run, interval_hours, enabled: true, last_run: None, run_count: 0, max_runs: None, created_at: now, metadata: HashMap::new(), }; self.tasks.insert(task_id.clone(), task); self.save_scheduler_data()?; Ok(task_id) } pub fn enable_task(&mut self, task_id: &str) -> Result<()> { if let Some(task) = self.tasks.get_mut(task_id) { task.enabled = true; self.save_scheduler_data()?; } Ok(()) } pub fn disable_task(&mut self, task_id: &str) -> Result<()> { if let Some(task) = self.tasks.get_mut(task_id) { task.enabled = false; self.save_scheduler_data()?; } Ok(()) } pub fn delete_task(&mut self, task_id: &str) -> Result<()> { self.tasks.remove(task_id); self.save_scheduler_data()?; Ok(()) } pub fn get_task(&self, task_id: &str) -> Option<&ScheduledTask> { self.tasks.get(task_id) } pub fn get_tasks(&self) -> &HashMap { &self.tasks } pub fn get_due_tasks(&self) -> Vec<&ScheduledTask> { let now = Utc::now(); self.tasks .values() .filter(|task| task.enabled && task.next_run <= now) .collect() } pub fn get_execution_history(&self, limit: Option) -> Vec<&TaskExecution> { let mut executions: Vec<_> = self.execution_history.iter().collect(); executions.sort_by(|a, b| b.execution_time.cmp(&a.execution_time)); match limit { Some(limit) => executions.into_iter().take(limit).collect(), None => executions, } } pub fn get_scheduler_stats(&self) -> SchedulerStats { let total_tasks = self.tasks.len(); let enabled_tasks = self.tasks.values().filter(|task| task.enabled).count(); let due_tasks = self.get_due_tasks().len(); let total_executions = self.execution_history.len(); let successful_executions = self.execution_history.iter() .filter(|exec| exec.success) .count(); let today = Utc::now().date_naive(); let today_executions = self.execution_history.iter() .filter(|exec| exec.execution_time.date_naive() == today) .count(); let avg_duration = if total_executions > 0 { self.execution_history.iter() .map(|exec| exec.duration_ms) .sum::() as f64 / total_executions as f64 } else { 0.0 }; SchedulerStats { total_tasks, enabled_tasks, due_tasks, total_executions, successful_executions, today_executions, success_rate: if total_executions > 0 { successful_executions as f64 / total_executions as f64 } else { 0.0 }, avg_duration_ms: avg_duration, } } fn create_default_tasks(&mut self) -> Result<()> { let now = Utc::now(); // Daily maintenance task - run every day at 3 AM let mut daily_maintenance_time = now.date_naive().and_hms_opt(3, 0, 0).unwrap().and_utc(); if daily_maintenance_time <= now { daily_maintenance_time = daily_maintenance_time + Duration::days(1); } self.create_task( TaskType::DailyMaintenance, daily_maintenance_time, Some(24), // 24 hours = 1 day )?; // Auto transmission check - every 4 hours self.create_task( TaskType::AutoTransmission, now + Duration::hours(1), Some(4), )?; // Breakthrough check - every 2 hours self.create_task( TaskType::BreakthroughCheck, now + Duration::minutes(30), Some(2), )?; // Maintenance transmission - once per day let mut maintenance_time = now.date_naive().and_hms_opt(12, 0, 0).unwrap().and_utc(); if maintenance_time <= now { maintenance_time = maintenance_time + Duration::days(1); } self.create_task( TaskType::MaintenanceTransmission, maintenance_time, Some(24), // 24 hours = 1 day )?; Ok(()) } fn load_scheduler_data(config: &Config) -> Result<(HashMap, Vec)> { let tasks_file = config.scheduler_tasks_file(); let history_file = config.scheduler_history_file(); let tasks = if tasks_file.exists() { let content = std::fs::read_to_string(tasks_file) .context("Failed to read scheduler tasks file")?; serde_json::from_str(&content) .context("Failed to parse scheduler tasks file")? } else { HashMap::new() }; let history = if history_file.exists() { let content = std::fs::read_to_string(history_file) .context("Failed to read scheduler history file")?; serde_json::from_str(&content) .context("Failed to parse scheduler history file")? } else { Vec::new() }; Ok((tasks, history)) } fn save_scheduler_data(&self) -> Result<()> { // Save tasks let tasks_content = serde_json::to_string_pretty(&self.tasks) .context("Failed to serialize scheduler tasks")?; std::fs::write(&self.config.scheduler_tasks_file(), tasks_content) .context("Failed to write scheduler tasks file")?; // Save execution history let history_content = serde_json::to_string_pretty(&self.execution_history) .context("Failed to serialize scheduler history")?; std::fs::write(&self.config.scheduler_history_file(), history_content) .context("Failed to write scheduler history file")?; Ok(()) } } // Type alias for compatibility with CLI interface pub type Scheduler = AIScheduler; impl Scheduler { pub fn list_tasks(&self) -> Result> { let tasks: Vec = self.tasks .values() .map(|task| ScheduledTaskInfo { name: task.task_type.to_string(), schedule: match task.interval_hours { Some(hours) => format!("Every {} hours", hours), None => "One-time".to_string(), }, next_run: task.next_run, enabled: task.enabled, }) .collect(); Ok(tasks) } } #[derive(Debug, Clone)] pub struct SchedulerStats { pub total_tasks: usize, pub enabled_tasks: usize, pub due_tasks: usize, pub total_executions: usize, pub successful_executions: usize, pub today_executions: usize, pub success_rate: f64, pub avg_duration_ms: f64, } #[derive(Debug, Clone)] pub struct ScheduledTaskInfo { pub name: String, pub schedule: String, pub next_run: DateTime, pub enabled: bool, }