commons/src/task.rs

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::agents::AgentType;
use crate::enums::{TaskResultType, TaskStatus, TestStrategy};
use crate::todo::ToDo;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Task {
    pub id: Uuid,
    pub todo_id: Uuid,
    pub description: String,
    pub status: TaskStatus,
    pub assigned_krill: Option<Uuid>,
    pub assigned_pod: Option<Uuid>,
    pub result: Option<TaskResult>,
    pub logs: Vec<String>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
    pub payload: Option<serde_json::Value>,
}

impl Default for Task {
    fn default() -> Self {
        Self {
            id: Uuid::new_v4(),
            todo_id: Uuid::nil(),
            description: String::new(),
            status: TaskStatus::default(),
            assigned_krill: None,
            assigned_pod: None,
            result: None,
            logs: Vec::new(),
            started_at: None,
            completed_at: None,
            payload: None,
        }
    }
}

impl Task {
    pub fn new(todo_id: Uuid, description: String) -> Self {
        Self {
            id: Uuid::new_v4(),
            todo_id,
            description,
            status: TaskStatus::default(),
            assigned_krill: None,
            assigned_pod: None,
            result: None,
            logs: Vec::new(),
            started_at: None,
            completed_at: None,
            payload: None,
        }
    }

    pub fn with_payload<T: Serialize>(mut self, payload: &T) -> Self {
        self.payload = serde_json::to_value(payload).ok();
        self
    }

    pub fn start(&mut self) -> Result<(), String> {
        match self.status {
            TaskStatus::Pending => {
                self.status = TaskStatus::Running;
                self.started_at = Some(Utc::now());
                Ok(())
            }
            _ => Err(format!("cannot start task from status: {:?}", self.status)),
        }
    }

    pub fn complete(&mut self, result: TaskResult) -> Result<(), String> {
        match self.status {
            TaskStatus::Running => {
                self.status = TaskStatus::Completed;
                self.completed_at = Some(Utc::now());
                self.result = Some(result);
                Ok(())
            }
            _ => Err(format!(
                "cannot complete task from status: {:?}",
                self.status
            )),
        }
    }

    pub fn fail(&mut self, error: String) -> Result<(), String> {
        match self.status {
            TaskStatus::Running => {
                self.status = TaskStatus::Failed;
                self.completed_at = Some(Utc::now());
                self.result = Some(TaskResult::Error(error));
                Ok(())
            }
            _ => Err(format!("cannot fail task from status: {:?}", self.status)),
        }
    }

    pub fn retry(&mut self) -> Result<(), String> {
        match self.status {
            TaskStatus::Failed | TaskStatus::Pending => {
                self.status = TaskStatus::Pending;
                self.started_at = None;
                self.completed_at = None;
                self.result = None;
                self.logs.clear();
                Ok(())
            }
            _ => Err(format!("cannot retry task from status: {:?}", self.status)),
        }
    }

    pub fn duration_ms(&self) -> Option<i64> {
        match (self.started_at, self.completed_at) {
            (Some(start), Some(end)) => Some((end - start).num_milliseconds()),
            _ => None,
        }
    }

    pub fn add_log(&mut self, log: String) {
        self.logs.push(log);
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum TaskResult {
    Success(String),
    Error(String),
    Split(Vec<ToDo>),
}

impl TaskResult {
    pub fn result_type(&self) -> TaskResultType {
        match self {
            TaskResult::Success(_) => TaskResultType::Success,
            TaskResult::Error(_) => TaskResultType::Error,
            TaskResult::Split(_) => TaskResultType::Split,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestCase {
    pub name: String,
    pub command: String,
    pub expected_output: Option<String>,
    #[serde(default)]
    pub test_strategy: Option<TestStrategy>,
    #[serde(default)]
    pub suggested_test_files: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceLimits {
    pub cpu_cores: f32,
    pub memory_mb: u64,
    pub gpu_memory_mb: Option<u64>,
    pub timeout_seconds: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionBudget {
    pub max_tool_calls: u32,
    pub max_tokens: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileRange {
    pub path: String,
    pub line_start: u32,
    pub line_end: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConnectorConfig {
    pub api_key: String,
    pub base_url: String,
    pub model: String,
    pub max_tokens: Option<i32>,
    pub temperature: Option<f32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentTaskSpec {
    pub task_id: Uuid,
    pub title: String,
    pub description: String,
    pub board_snapshot: serde_json::Value,
    pub assembled_context_id: Uuid,
    pub selected_files: Vec<FileRange>,
    pub entry_points: Vec<String>,
    pub tests: Option<Vec<TestCase>>,
    pub success_criteria: Option<String>,
    pub token_budget: usize,
    pub resource_limits: ResourceLimits,
    pub parent_task_id: Option<Uuid>,
    pub pre_approved_plan_id: Option<Uuid>,
    pub execution_token: Option<Uuid>,
    pub session_budget: Option<SessionBudget>,
    pub agent_type: Option<AgentType>,
    #[serde(default)]
    pub suggested_test_files: Vec<String>,
    #[serde(default)]
    pub project_id: Option<Uuid>,
    #[serde(default)]
    pub connector_config: Option<ConnectorConfig>,
    #[serde(default)]
    pub tools_json: Option<String>,
    /* Overrides the agent-type default system prompt when set. Takes highest
       precedence in the resolution chain (above DB overrides and defaults). */
    #[serde(default)]
    pub system_prompt: Option<String>,
}

impl AgentTaskSpec {
    pub fn new(task_id: Uuid, title: String, assembled_context_id: Uuid) -> Self {
        Self {
            task_id,
            title,
            description: String::new(),
            board_snapshot: serde_json::json!({}),
            assembled_context_id,
            selected_files: Vec::new(),
            entry_points: Vec::new(),
            tests: None,
            success_criteria: None,
            token_budget: 1000000, // Default 1M tokens
            resource_limits: ResourceLimits {
                cpu_cores: 1.0,
                memory_mb: 2048,
                gpu_memory_mb: None,
                timeout_seconds: 3600, // 1 hour
            },
            parent_task_id: None,
            pre_approved_plan_id: None,
            execution_token: None,
            session_budget: None,
            agent_type: None,
            suggested_test_files: Vec::new(),
            project_id: None,
            connector_config: None,
            tools_json: None,
            system_prompt: None,
        }
    }
}

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

    #[test]
    fn test_task_new() {
        let todo_id = Uuid::new_v4();
        let task = Task::new(todo_id, "Test task description".to_string());
        assert!(!task.id.is_nil());
        assert_eq!(task.todo_id, todo_id);
        assert_eq!(task.description, "Test task description");
        assert_eq!(task.status, TaskStatus::Pending);
    }

    #[test]
    fn test_task_start() {
        let mut task = Task::default();
        assert!(task.start().is_ok());
        assert_eq!(task.status, TaskStatus::Running);
        assert!(task.started_at.is_some());
    }

    #[test]
    fn test_task_complete() {
        let mut task = Task::default();
        task.status = TaskStatus::Running;
        let result = TaskResult::Success("output".to_string());
        assert!(task.complete(result.clone()).is_ok());
        assert_eq!(task.status, TaskStatus::Completed);
        assert!(task.completed_at.is_some());
        assert_eq!(task.result, Some(result));
    }

    #[test]
    fn test_task_duration_ms() {
        let mut task = Task::default();
        task.start().expect("task should start");
        std::thread::sleep(std::time::Duration::from_millis(10));
        task.complete(TaskResult::Success("done".to_string()))
            .expect("task should complete");

        let duration = task.duration_ms();
        assert!(duration.is_some());
        assert!(duration.unwrap() >= 10);
    }

    #[test]
    fn test_task_retry() {
        let mut task = Task::default();
        task.status = TaskStatus::Failed;
        task.result = Some(TaskResult::Error("failed".to_string()));
        assert!(task.retry().is_ok());
        assert_eq!(task.status, TaskStatus::Pending);
        assert!(task.result.is_none());
    }

    #[test]
    fn test_task_default_values() {
        let task = Task::default();
        assert_eq!(task.status, TaskStatus::Pending);
        assert!(task.result.is_none());
        assert!(task.logs.is_empty());
        assert!(task.started_at.is_none());
        assert!(task.completed_at.is_none());
    }

    #[test]
    fn test_task_serialization_roundtrip() {
        let original = Task::new(Uuid::new_v4(), "Test task".to_string());
        let serialized = serde_json::to_string(&original).unwrap();
        let deserialized: Task = serde_json::from_str(&serialized).unwrap();

        assert_eq!(original.id, deserialized.id);
        assert_eq!(original.todo_id, deserialized.todo_id);
        assert_eq!(original.description, deserialized.description);
        assert_eq!(original.status, deserialized.status);
    }

    #[test]
    fn test_task_fail() {
        let mut task = Task::default();
        task.status = TaskStatus::Running;
        assert!(task.fail("error message".to_string()).is_ok());
        assert_eq!(task.status, TaskStatus::Failed);
        assert!(task.result.is_some());
    }

    #[test]
    fn test_task_invalid_transition() {
        let mut task = Task::default();
        // Cannot complete a pending task directly
        let result = TaskResult::Success("done".to_string());
        assert!(task.complete(result).is_err());
    }
}