executor.rs

//! Background agent executor and manager.
//!
//! Provides the core infrastructure for spawning and managing background agents
//! as tokio tasks with proper lifecycle management.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
use tokio::sync::{broadcast, mpsc, oneshot, RwLock};
use tokio::task::JoinHandle;
use uuid::Uuid;

use super::events::AgentEvent;
use super::messaging::{AgentMailbox, MessageRouter};

/// Default maximum concurrent background agents.
pub const DEFAULT_MAX_CONCURRENT: usize = 5;

/// Default timeout for background agents (30 minutes).
pub const DEFAULT_TIMEOUT_SECS: u64 = 1800;

/// Grace period for cancellation (5 seconds).
pub const CANCEL_GRACE_PERIOD_SECS: u64 = 5;

/// Errors from background agent operations.
#[derive(Error, Debug)]
pub enum BackgroundAgentError {
    /// Too many concurrent agents.
    #[error("Too many concurrent agents (max: {max})")]
    TooManyAgents { max: usize },

    /// Agent not found.
    #[error("Agent not found: {0}")]
    AgentNotFound(String),

    /// Agent already completed.
    #[error("Agent already completed: {0}")]
    AlreadyCompleted(String),

    /// Timeout waiting for agent.
    #[error("Timeout waiting for agent: {0}")]
    Timeout(String),

    /// Agent execution failed.
    #[error("Agent execution failed: {0}")]
    ExecutionFailed(String),

    /// Internal error.
    #[error("Internal error: {0}")]
    Internal(String),

    /// Cancellation failed.
    #[error("Failed to cancel agent: {0}")]
    CancelFailed(String),
}

/// Result type for background agent operations.
pub type Result<T> = std::result::Result<T, BackgroundAgentError>;

/// Status of a background agent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AgentStatus {
    /// Agent is initializing.
    Initializing,
    /// Agent is running.
    Running,
    /// Agent completed successfully.
    Completed,
    /// Agent failed with error.
    Failed(String),
    /// Agent was cancelled.
    Cancelled,
}

impl AgentStatus {
    /// Check if the status is terminal.
    pub fn is_terminal(&self) -> bool {
        matches!(
            self,
            AgentStatus::Completed | AgentStatus::Failed(_) | AgentStatus::Cancelled
        )
    }

    /// Get a display string for the status.
    pub fn as_str(&self) -> &'static str {
        match self {
            AgentStatus::Initializing => "initializing",
            AgentStatus::Running => "running",
            AgentStatus::Completed => "completed",
            AgentStatus::Failed(_) => "failed",
            AgentStatus::Cancelled => "cancelled",
        }
    }
}

impl Default for AgentStatus {
    fn default() -> Self {
        AgentStatus::Initializing
    }
}

/// Result of a background agent execution.
#[derive(Debug, Clone)]
pub struct AgentResult {
    /// Summary of what the agent accomplished.
    pub summary: String,
    /// Output data (if any).
    pub output: Option<serde_json::Value>,
    /// Tokens used.
    pub tokens_used: u64,
    /// Files modified (paths).
    pub files_modified: Vec<String>,
    /// Whether the execution was successful.
    pub success: bool,
    /// Error message (if failed).
    pub error: Option<String>,
}

impl Default for AgentResult {
    fn default() -> Self {
        Self {
            summary: String::new(),
            output: None,
            tokens_used: 0,
            files_modified: Vec::new(),
            success: true,
            error: None,
        }
    }
}

impl AgentResult {
    /// Create a successful result.
    pub fn success(summary: impl Into<String>) -> Self {
        Self {
            summary: summary.into(),
            success: true,
            ..Default::default()
        }
    }

    /// Create a failed result.
    pub fn failure(error: impl Into<String>) -> Self {
        let error = error.into();
        Self {
            summary: format!("Failed: {}", &error),
            success: false,
            error: Some(error),
            ..Default::default()
        }
    }

    /// Set the tokens used.
    pub fn with_tokens(mut self, tokens: u64) -> Self {
        self.tokens_used = tokens;
        self
    }

    /// Set the files modified.
    pub fn with_files(mut self, files: Vec<String>) -> Self {
        self.files_modified = files;
        self
    }

    /// Set additional output data.
    pub fn with_output(mut self, output: serde_json::Value) -> Self {
        self.output = Some(output);
        self
    }
}

/// Configuration for spawning a background agent.
#[derive(Debug, Clone)]
pub struct AgentConfig {
    /// Task description/prompt.
    pub task: String,
    /// Agent type (e.g., "general", "explore", "research").
    pub agent_type: String,
    /// Custom system prompt (optional).
    pub system_prompt: Option<String>,
    /// Maximum execution time in seconds.
    pub timeout_secs: u64,
    /// Priority (higher = more important).
    pub priority: i32,
    /// Tags for categorization.
    pub tags: Vec<String>,
    /// Whether this agent can modify files.
    pub can_modify: bool,
    /// Model override.
    pub model: Option<String>,
    /// Temperature override.
    pub temperature: Option<f32>,
}

impl AgentConfig {
    /// Create a new agent config with default settings.
    pub fn new(task: impl Into<String>) -> Self {
        Self {
            task: task.into(),
            agent_type: "general".to_string(),
            system_prompt: None,
            timeout_secs: DEFAULT_TIMEOUT_SECS,
            priority: 0,
            tags: Vec::new(),
            can_modify: false,
            model: None,
            temperature: None,
        }
    }

    /// Create a config for background execution.
    pub fn background(task: impl Into<String>) -> Self {
        Self::new(task).with_type("background")
    }

    /// Create a test config.
    #[cfg(test)]
    pub fn test() -> Self {
        Self::new("test task").with_timeout(10)
    }

    /// Create a long-running test config.
    #[cfg(test)]
    pub fn long_running() -> Self {
        Self::new("long running task").with_timeout(300)
    }

    /// Set the agent type.
    pub fn with_type(mut self, agent_type: impl Into<String>) -> Self {
        self.agent_type = agent_type.into();
        self
    }

    /// Set the timeout.
    pub fn with_timeout(mut self, secs: u64) -> Self {
        self.timeout_secs = secs;
        self
    }

    /// Set the priority.
    pub fn with_priority(mut self, priority: i32) -> Self {
        self.priority = priority;
        self
    }

    /// Add a tag.
    pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
        self.tags.push(tag.into());
        self
    }

    /// Allow file modifications.
    pub fn with_modify_access(mut self) -> Self {
        self.can_modify = true;
        self
    }

    /// Set the model.
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// Set the temperature.
    pub fn with_temperature(mut self, temp: f32) -> Self {
        self.temperature = Some(temp);
        self
    }

    /// Set the system prompt.
    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(prompt.into());
        self
    }
}

/// A background agent instance.
pub struct BackgroundAgent {
    /// Unique agent ID.
    pub id: String,
    /// Task description.
    pub task: String,
    /// Agent type.
    pub agent_type: String,
    /// Current status.
    status: AgentStatus,
    /// Task handle.
    handle: Option<JoinHandle<AgentResult>>,
    /// Status update sender.
    status_tx: mpsc::Sender<AgentStatus>,
    /// Status update receiver.
    status_rx: mpsc::Receiver<AgentStatus>,
    /// Cancellation sender.
    cancel_tx: Option<oneshot::Sender<()>>,
    /// Agent mailbox for inter-agent communication.
    pub mailbox: AgentMailbox,
    /// Start time.
    pub started_at: Instant,
    /// Completion result (set when agent completes).
    result: Option<AgentResult>,
    /// Tokens used so far.
    pub tokens_used: u64,
    /// Priority level.
    pub priority: i32,
    /// Associated tags.
    pub tags: Vec<String>,
}

impl BackgroundAgent {
    /// Create a new background agent.
    fn new(id: String, config: &AgentConfig, mailbox: AgentMailbox) -> Self {
        let (status_tx, status_rx) = mpsc::channel(32);
        Self {
            id,
            task: config.task.clone(),
            agent_type: config.agent_type.clone(),
            status: AgentStatus::Initializing,
            handle: None,
            status_tx,
            status_rx,
            cancel_tx: None,
            mailbox,
            started_at: Instant::now(),
            result: None,
            tokens_used: 0,
            priority: config.priority,
            tags: config.tags.clone(),
        }
    }

    /// Get the current status.
    pub fn status(&self) -> &AgentStatus {
        &self.status
    }

    /// Get the duration since start.
    pub fn duration(&self) -> Duration {
        self.started_at.elapsed()
    }

    /// Get the duration in a human-readable format.
    pub fn duration_display(&self) -> String {
        let secs = self.duration().as_secs();
        if secs < 60 {
            format!("{}s", secs)
        } else if secs < 3600 {
            format!("{}m {}s", secs / 60, secs % 60)
        } else {
            format!("{}h {}m", secs / 3600, (secs % 3600) / 60)
        }
    }

    /// Check if the agent is still running.
    pub fn is_running(&self) -> bool {
        matches!(
            self.status,
            AgentStatus::Running | AgentStatus::Initializing
        )
    }

    /// Get the result (if completed).
    pub fn result(&self) -> Option<&AgentResult> {
        self.result.as_ref()
    }

    /// Get a status snapshot for display.
    pub fn snapshot(&self) -> AgentSnapshot {
        AgentSnapshot {
            id: self.id.clone(),
            task: self.task.clone(),
            agent_type: self.agent_type.clone(),
            status: self.status.clone(),
            duration: self.duration(),
            tokens_used: self.tokens_used,
            priority: self.priority,
        }
    }
}

/// A snapshot of agent state for display.
#[derive(Debug, Clone)]
pub struct AgentSnapshot {
    /// Agent ID.
    pub id: String,
    /// Task description.
    pub task: String,
    /// Agent type.
    pub agent_type: String,
    /// Current status.
    pub status: AgentStatus,
    /// Duration since start.
    pub duration: Duration,
    /// Tokens used.
    pub tokens_used: u64,
    /// Priority.
    pub priority: i32,
}

/// Manager for background agents.
pub struct BackgroundAgentManager {
    /// Active agents by ID.
    agents: Arc<RwLock<HashMap<String, BackgroundAgent>>>,
    /// Event broadcaster.
    event_tx: broadcast::Sender<AgentEvent>,
    /// Message router for inter-agent communication.
    message_router: Arc<MessageRouter>,
    /// Maximum concurrent agents.
    max_concurrent: usize,
}

impl BackgroundAgentManager {
    /// Create a new background agent manager.
    pub fn new(max_concurrent: usize) -> Self {
        let (event_tx, _) = broadcast::channel(256);
        Self {
            agents: Arc::new(RwLock::new(HashMap::new())),
            event_tx,
            message_router: Arc::new(MessageRouter::new()),
            max_concurrent,
        }
    }

    /// Create with default settings.
    pub fn default_manager() -> Self {
        Self::new(DEFAULT_MAX_CONCURRENT)
    }

    /// Spawn a new background agent.
    pub async fn spawn(&self, config: AgentConfig) -> Result<String> {
        // Check concurrent limit
        let agents = self.agents.read().await;
        let active_count = agents.values().filter(|a| a.is_running()).count();
        if active_count >= self.max_concurrent {
            return Err(BackgroundAgentError::TooManyAgents {
                max: self.max_concurrent,
            });
        }
        drop(agents);

        // Generate ID
        let id = Uuid::new_v4().to_string();

        // Create mailbox
        let mailbox = self.message_router.create_mailbox(&id);

        // Create agent
        let mut agent = BackgroundAgent::new(id.clone(), &config, mailbox);

        // Create cancellation channel
        let (cancel_tx, cancel_rx) = oneshot::channel();
        agent.cancel_tx = Some(cancel_tx);

        // Create task
        let event_tx = self.event_tx.clone();
        let status_tx = agent.status_tx.clone();
        let task_description = config.task.clone();
        let timeout_secs = config.timeout_secs;
        let agent_id = id.clone();

        let handle = tokio::spawn(async move {
            run_background_agent(
                agent_id,
                task_description,
                timeout_secs,
                cancel_rx,
                status_tx,
                event_tx,
            )
            .await
        });

        agent.handle = Some(handle);
        agent.status = AgentStatus::Running;

        // Broadcast started event
        let _ = self.event_tx.send(AgentEvent::Started {
            id: id.clone(),
            task: config.task.clone(),
            timestamp_ms: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_millis() as u64,
        });

        // Insert into registry
        let mut agents = self.agents.write().await;
        agents.insert(id.clone(), agent);

        tracing::info!(agent_id = %id, "Spawned background agent");

        Ok(id)
    }

    /// List all agents (active and completed).
    pub async fn list(&self) -> Vec<AgentSnapshot> {
        let agents = self.agents.read().await;
        agents.values().map(|a| a.snapshot()).collect()
    }

    /// List only active agents.
    pub async fn list_active(&self) -> Vec<AgentSnapshot> {
        let agents = self.agents.read().await;
        agents
            .values()
            .filter(|a| a.is_running())
            .map(|a| a.snapshot())
            .collect()
    }

    /// Get the status of a specific agent.
    pub async fn get_status(&self, id: &str) -> Option<AgentStatus> {
        let agents = self.agents.read().await;
        agents.get(id).map(|a| a.status.clone())
    }

    /// Get a snapshot of a specific agent.
    pub async fn get_agent(&self, id: &str) -> Option<AgentSnapshot> {
        let agents = self.agents.read().await;
        agents.get(id).map(|a| a.snapshot())
    }

    /// Wait for an agent to complete with timeout.
    pub async fn wait(&self, id: &str, timeout: Duration) -> Result<AgentResult> {
        let deadline = Instant::now() + timeout;

        loop {
            // Check if agent exists and get its status
            let agents = self.agents.read().await;
            let agent = agents
                .get(id)
                .ok_or_else(|| BackgroundAgentError::AgentNotFound(id.to_string()))?;

            if agent.status.is_terminal() {
                return agent.result.clone().ok_or_else(|| {
                    BackgroundAgentError::Internal("No result available".to_string())
                });
            }

            drop(agents);

            // Check timeout
            if Instant::now() >= deadline {
                return Err(BackgroundAgentError::Timeout(id.to_string()));
            }

            // Wait a bit before checking again
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    /// Cancel an agent.
    pub async fn cancel(&self, id: &str) -> Result<()> {
        let mut agents = self.agents.write().await;
        let agent = agents
            .get_mut(id)
            .ok_or_else(|| BackgroundAgentError::AgentNotFound(id.to_string()))?;

        if agent.status.is_terminal() {
            return Err(BackgroundAgentError::AlreadyCompleted(id.to_string()));
        }

        // Send cancellation signal
        if let Some(cancel_tx) = agent.cancel_tx.take() {
            let _ = cancel_tx.send(());
        }

        // Update status
        agent.status = AgentStatus::Cancelled;
        agent.result = Some(AgentResult::failure("Cancelled by user"));

        // Broadcast event
        let _ = self.event_tx.send(AgentEvent::Cancelled {
            id: id.to_string(),
            reason: Some("User cancelled".to_string()),
        });

        tracing::info!(agent_id = %id, "Cancelled background agent");

        Ok(())
    }

    /// Cancel all running agents.
    pub async fn cancel_all(&self) -> Vec<Result<()>> {
        let agents = self.agents.read().await;
        let running_ids: Vec<String> = agents
            .values()
            .filter(|a| a.is_running())
            .map(|a| a.id.clone())
            .collect();
        drop(agents);

        let mut results = Vec::with_capacity(running_ids.len());
        for id in running_ids {
            results.push(self.cancel(&id).await);
        }
        results
    }

    /// Subscribe to agent events.
    pub fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
        self.event_tx.subscribe()
    }

    /// Get the event sender for external use.
    pub fn event_sender(&self) -> broadcast::Sender<AgentEvent> {
        self.event_tx.clone()
    }

    /// Get the number of active agents.
    pub async fn active_count(&self) -> usize {
        let agents = self.agents.read().await;
        agents.values().filter(|a| a.is_running()).count()
    }

    /// Get the total number of agents (including completed).
    pub async fn total_count(&self) -> usize {
        let agents = self.agents.read().await;
        agents.len()
    }

    /// Clean up completed agents older than the given age.
    pub async fn cleanup(&self, max_age: Duration) {
        let mut agents = self.agents.write().await;
        let now = Instant::now();
        agents.retain(|_, agent| {
            if agent.status.is_terminal() {
                agent.duration() < max_age
            } else {
                true
            }
        });
        let _ = now; // Suppress unused warning
    }

    /// Get the message router for inter-agent communication.
    pub fn message_router(&self) -> &Arc<MessageRouter> {
        &self.message_router
    }

    /// Update the status of an agent (for internal use).
    pub(crate) async fn update_status(
        &self,
        id: &str,
        status: AgentStatus,
        result: Option<AgentResult>,
    ) {
        let mut agents = self.agents.write().await;
        if let Some(agent) = agents.get_mut(id) {
            agent.status = status;
            if let Some(r) = result {
                agent.result = Some(r);
            }
        }
    }
}

impl Default for BackgroundAgentManager {
    fn default() -> Self {
        Self::default_manager()
    }
}

impl Drop for BackgroundAgentManager {
    fn drop(&mut self) {
        // Cancel all agents synchronously (best effort)
        // Note: This is a sync drop, so we can't await
        // The tokio tasks will be aborted when dropped
        tracing::debug!("BackgroundAgentManager dropped, agents will be aborted");
    }
}

/// Run a background agent task.
async fn run_background_agent(
    id: String,
    task: String,
    timeout_secs: u64,
    mut cancel_rx: oneshot::Receiver<()>,
    status_tx: mpsc::Sender<AgentStatus>,
    event_tx: broadcast::Sender<AgentEvent>,
) -> AgentResult {
    let start = Instant::now();

    // Update status to running
    let _ = status_tx.send(AgentStatus::Running).await;

    // Simulate agent work with timeout and cancellation
    let timeout_duration = Duration::from_secs(timeout_secs);

    tokio::select! {
        // Timeout
        _ = tokio::time::sleep(timeout_duration) => {
            let _ = status_tx.send(AgentStatus::Failed("Timeout".to_string())).await;
            let _ = event_tx.send(AgentEvent::Failed {
                id: id.clone(),
                error: "Agent timed out".to_string(),
                duration_ms: start.elapsed().as_millis() as u64,
                recoverable: false,
            });
            AgentResult::failure("Agent execution timed out")
        }

        // Cancellation
        _ = &mut cancel_rx => {
            let _ = status_tx.send(AgentStatus::Cancelled).await;
            let _ = event_tx.send(AgentEvent::Cancelled {
                id: id.clone(),
                reason: Some("Cancelled".to_string()),
            });
            AgentResult::failure("Cancelled")
        }

        // Simulated work completion (in real implementation, this would be actual agent work)
        result = simulate_agent_work(&id, &task, &event_tx) => {
            let duration_ms = start.elapsed().as_millis() as u64;

            if result.success {
                let _ = status_tx.send(AgentStatus::Completed).await;
                let _ = event_tx.send(AgentEvent::Completed {
                    id: id.clone(),
                    summary: result.summary.clone(),
                    duration_ms,
                    tokens_used: Some(result.tokens_used),
                });
            } else {
                let _ = status_tx.send(AgentStatus::Failed(result.error.clone().unwrap_or_default())).await;
                let _ = event_tx.send(AgentEvent::Failed {
                    id: id.clone(),
                    error: result.error.clone().unwrap_or_else(|| "Unknown error".to_string()),
                    duration_ms,
                    recoverable: false,
                });
            }

            result
        }
    }
}

/// Simulate agent work (placeholder for actual agent execution).
async fn simulate_agent_work(
    id: &str,
    task: &str,
    event_tx: &broadcast::Sender<AgentEvent>,
) -> AgentResult {
    // Send initial progress
    let _ = event_tx.send(AgentEvent::Progress {
        id: id.to_string(),
        message: "Starting task execution".to_string(),
        step: Some(1),
        total_steps: Some(3),
    });

    // Simulate some work
    tokio::time::sleep(Duration::from_millis(100)).await;

    let _ = event_tx.send(AgentEvent::Progress {
        id: id.to_string(),
        message: "Processing".to_string(),
        step: Some(2),
        total_steps: Some(3),
    });

    tokio::time::sleep(Duration::from_millis(100)).await;

    let _ = event_tx.send(AgentEvent::Progress {
        id: id.to_string(),
        message: "Finalizing".to_string(),
        step: Some(3),
        total_steps: Some(3),
    });

    // Return successful result
    AgentResult::success(format!("Completed task: {}", truncate_str(task, 50))).with_tokens(150)
}

/// Truncate a string for display.
fn truncate_str(s: &str, max_len: usize) -> &str {
    if s.len() <= max_len {
        s
    } else {
        &s[..max_len]
    }
}

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

    #[test]
    fn test_agent_config_builder() {
        let config = AgentConfig::new("Test task")
            .with_type("explore")
            .with_timeout(60)
            .with_priority(5)
            .with_tag("important");

        assert_eq!(config.task, "Test task");
        assert_eq!(config.agent_type, "explore");
        assert_eq!(config.timeout_secs, 60);
        assert_eq!(config.priority, 5);
        assert!(config.tags.contains(&"important".to_string()));
    }

    #[test]
    fn test_agent_status_is_terminal() {
        assert!(!AgentStatus::Initializing.is_terminal());
        assert!(!AgentStatus::Running.is_terminal());
        assert!(AgentStatus::Completed.is_terminal());
        assert!(AgentStatus::Failed("error".to_string()).is_terminal());
        assert!(AgentStatus::Cancelled.is_terminal());
    }

    #[test]
    fn test_agent_result_builders() {
        let success = AgentResult::success("Done");
        assert!(success.success);
        assert_eq!(success.summary, "Done");

        let failure = AgentResult::failure("Error occurred");
        assert!(!failure.success);
        assert_eq!(failure.error, Some("Error occurred".to_string()));
    }

    #[tokio::test]
    async fn test_spawn_background_agent() {
        let manager = BackgroundAgentManager::new(5);
        let config = AgentConfig::test();

        let id = manager.spawn(config).await.unwrap();
        assert!(!id.is_empty());

        // Check it's in the list
        let agents = manager.list().await;
        assert_eq!(agents.len(), 1);
        assert_eq!(agents[0].id, id);
    }

    #[tokio::test]
    async fn test_max_concurrent_limit() {
        let manager = BackgroundAgentManager::new(2);

        // Spawn 2 agents
        manager.spawn(AgentConfig::long_running()).await.unwrap();
        manager.spawn(AgentConfig::long_running()).await.unwrap();

        // Third should fail
        let result = manager.spawn(AgentConfig::test()).await;
        assert!(matches!(
            result,
            Err(BackgroundAgentError::TooManyAgents { .. })
        ));
    }

    #[tokio::test]
    async fn test_agent_cancellation() {
        let manager = BackgroundAgentManager::new(5);
        let id = manager.spawn(AgentConfig::long_running()).await.unwrap();

        // Wait a bit for agent to start
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Cancel the agent
        manager.cancel(&id).await.unwrap();

        // Check status
        let status = manager.get_status(&id).await;
        assert!(matches!(status, Some(AgentStatus::Cancelled)));
    }

    #[tokio::test]
    async fn test_event_subscription() {
        let manager = BackgroundAgentManager::new(5);
        let mut events = manager.subscribe();

        // Spawn an agent
        let id = manager.spawn(AgentConfig::test()).await.unwrap();

        // Should receive started event
        let event = tokio::time::timeout(Duration::from_secs(1), events.recv())
            .await
            .unwrap()
            .unwrap();

        assert!(matches!(event, AgentEvent::Started { .. }));
        if let AgentEvent::Started { id: event_id, .. } = event {
            assert_eq!(event_id, id);
        }
    }

    #[test]
    fn test_truncate_str() {
        assert_eq!(truncate_str("short", 10), "short");
        assert_eq!(truncate_str("this is a long string", 10), "this is a ");
    }
}