vtcode-core 0.97.1

//! Optimized agent execution loop with state machine and predictive optimization

use anyhow::{Result, anyhow};
use hashbrown::HashMap;
use serde_json::Value;
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{RwLock, mpsc};
use tracing::{debug, error, info, warn};

use crate::core::memory_pool::global_pool;
use crate::tools::ToolCallRequest;
use crate::tools::async_pipeline::{
    AsyncToolPipeline, ExecutionContext, ExecutionPriority, ToolRequest,
};

/// Agent execution states
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AgentState {
    Idle,
    ProcessingPrompt,
    ExecutingTools,
    GeneratingResponse,
    WaitingForUser,
    Error { error_type: String },
    Shutdown,
}

/// Agent execution context with optimization hints
#[derive(Debug, Clone)]
pub struct AgentContext {
    pub session_id: String,
    pub current_state: AgentState,
    pub conversation_history: Vec<ConversationTurn>,
    pub active_tools: HashMap<String, ToolExecutionState>,
    pub performance_hints: PerformanceHints,
    pub resource_limits: ResourceLimits,
}

/// Single conversation turn
#[derive(Debug, Clone)]
pub struct ConversationTurn {
    pub id: String,
    pub user_message: String,
    pub agent_response: Option<String>,
    pub tool_calls: Vec<ToolCall>,
    pub timestamp: Instant,
    pub execution_time: Option<Duration>,
}

/// Tool execution state tracking
#[derive(Debug, Clone)]
pub struct ToolExecutionState {
    pub tool_name: String,
    pub status: ToolStatus,
    pub started_at: Instant,
    pub estimated_completion: Option<Instant>,
    pub resource_usage: ResourceUsage,
}

/// Tool execution status
#[derive(Debug, Clone, PartialEq)]
pub enum ToolStatus {
    Queued,
    Running,
    Completed,
    Failed { error: String },
    Cancelled,
}

/// Tool call information
#[derive(Debug, Clone)]
pub struct ToolCall {
    pub id: String,
    pub tool_name: String,
    pub args: Value,
    pub result: Option<Value>,
    pub execution_time: Option<Duration>,
}

/// Performance optimization hints
#[derive(Debug, Clone)]
pub struct PerformanceHints {
    pub predicted_tool_sequence: Vec<String>,
    pub cache_warming_candidates: Vec<String>,
    pub parallel_execution_groups: Vec<Vec<String>>,
    pub resource_intensive_tools: Vec<String>,
}

/// Resource usage tracking
#[derive(Debug, Clone, Default)]
pub struct ResourceUsage {
    pub cpu_percent: f64,
    pub memory_mb: u64,
    pub network_bytes: u64,
    pub disk_io_bytes: u64,
}

/// Resource limits for execution
#[derive(Debug, Clone)]
pub struct ResourceLimits {
    pub max_concurrent_tools: usize,
    pub max_memory_mb: u64,
    pub max_execution_time: Duration,
    pub max_tool_retries: usize,
}

/// State transition event
#[derive(Debug, Clone)]
pub struct StateTransition {
    pub from_state: AgentState,
    pub to_state: AgentState,
    pub trigger: TransitionTrigger,
    pub timestamp: Instant,
}

/// What triggered a state transition
#[derive(Debug, Clone)]
pub enum TransitionTrigger {
    UserInput,
    ToolCompletion,
    ToolFailure,
    Timeout,
    ResourceLimit,
    InternalError,
}

/// Optimized agent execution engine
pub struct OptimizedAgentEngine {
    /// Current agent context
    context: Arc<RwLock<AgentContext>>,

    /// Tool execution pipeline
    tool_pipeline: Arc<AsyncToolPipeline>,

    /// State transition history for learning
    state_history: Arc<RwLock<VecDeque<StateTransition>>>,

    /// Performance predictor
    predictor: Arc<PerformancePredictor>,

    /// Event channel for state changes
    state_tx: mpsc::UnboundedSender<StateTransition>,
    state_rx: Arc<RwLock<Option<mpsc::UnboundedReceiver<StateTransition>>>>,
}

/// Performance predictor for optimization
#[derive(Default)]
pub struct PerformancePredictor {
    /// Historical execution patterns
    execution_patterns: Arc<RwLock<HashMap<String, ExecutionPattern>>>,
}

/// Historical execution pattern
#[derive(Debug, Clone)]
pub struct ExecutionPattern {
    pub tool_sequence: Vec<String>,
    pub avg_execution_time: Duration,
    pub success_rate: f64,
    pub resource_usage: ResourceUsage,
    pub frequency: u64,
}

impl OptimizedAgentEngine {
    pub fn new(session_id: String, tool_pipeline: Arc<AsyncToolPipeline>) -> Self {
        let (state_tx, state_rx) = mpsc::unbounded_channel();

        let context = AgentContext {
            session_id,
            current_state: AgentState::Idle,
            conversation_history: Vec::new(),
            active_tools: HashMap::new(),
            performance_hints: PerformanceHints {
                predicted_tool_sequence: Vec::new(),
                cache_warming_candidates: Vec::new(),
                parallel_execution_groups: Vec::new(),
                resource_intensive_tools: Vec::new(),
            },
            resource_limits: ResourceLimits {
                max_concurrent_tools: 4,
                max_memory_mb: 1024,
                max_execution_time: Duration::from_secs(300),
                max_tool_retries: 3,
            },
        };

        Self {
            context: Arc::new(RwLock::new(context)),
            tool_pipeline,
            state_history: Arc::new(RwLock::new(VecDeque::with_capacity(1000))),
            predictor: Arc::new(PerformancePredictor::new()),
            state_tx,
            state_rx: Arc::new(RwLock::new(Some(state_rx))),
        }
    }

    /// Start the optimized agent execution loop
    pub async fn start(&mut self) -> Result<()> {
        info!("Starting optimized agent engine");

        // Start background state monitoring
        self.start_state_monitor().await?;

        // Start performance prediction updates
        self.start_performance_predictor().await?;

        // Main execution loop
        loop {
            let current_state = self.context.read().await.current_state.clone();

            match current_state {
                AgentState::Idle => {
                    self.handle_idle_state().await?;
                }
                AgentState::ProcessingPrompt => {
                    self.handle_processing_state().await?;
                }
                AgentState::ExecutingTools => {
                    self.handle_tool_execution_state().await?;
                }
                AgentState::GeneratingResponse => {
                    self.handle_response_generation_state().await?;
                }
                AgentState::WaitingForUser => {
                    self.handle_waiting_state().await?;
                }
                AgentState::Error { error_type } => {
                    self.handle_error_state(&error_type).await?;
                }
                AgentState::Shutdown => {
                    info!("Agent shutdown requested");
                    break;
                }
            }

            // Small delay to prevent busy waiting
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        Ok(())
    }

    /// Handle idle state - wait for user input
    async fn handle_idle_state(&self) -> Result<()> {
        debug!("Agent in idle state, waiting for input");

        // Perform background optimizations while idle
        self.optimize_while_idle().await?;

        // Transition to processing when input is available
        // (This would be triggered by external input in a real implementation)
        tokio::time::sleep(Duration::from_millis(100)).await;

        Ok(())
    }

    /// Handle prompt processing state
    async fn handle_processing_state(&self) -> Result<()> {
        debug!("Processing user prompt");

        let start_time = Instant::now();

        // Get performance predictions for this prompt
        let predictions = self
            .predictor
            .predict_execution_pattern("user_prompt")
            .await;

        // Pre-warm caches based on predictions
        if let Some(pattern) = predictions {
            self.pre_warm_caches(&pattern.tool_sequence).await?;
        }

        // Simulate prompt processing
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Transition to tool execution if tools are needed
        self.transition_state(
            AgentState::ProcessingPrompt,
            AgentState::ExecutingTools,
            TransitionTrigger::ToolCompletion,
        )
        .await?;

        let processing_time = start_time.elapsed();
        debug!("Prompt processing completed in {:?}", processing_time);

        Ok(())
    }

    /// Handle tool execution state with optimization
    async fn handle_tool_execution_state(&self) -> Result<()> {
        debug!("Executing tools with optimization");

        let parallel_execution_groups = {
            let context = self.context.read().await;
            context.performance_hints.parallel_execution_groups.clone()
        };

        // Execute tools in parallel groups when possible
        for group in &parallel_execution_groups {
            self.execute_tool_group_parallel(group).await?;
        }

        // Check if all tools are complete
        let all_complete = {
            let context = self.context.read().await;
            context.active_tools.values().all(|state| {
                matches!(
                    state.status,
                    ToolStatus::Completed | ToolStatus::Failed { .. }
                )
            })
        };

        if all_complete {
            self.transition_state(
                AgentState::ExecutingTools,
                AgentState::GeneratingResponse,
                TransitionTrigger::ToolCompletion,
            )
            .await?;
        }

        Ok(())
    }

    /// Handle response generation state
    async fn handle_response_generation_state(&self) -> Result<()> {
        debug!("Generating response");

        // Use optimized LLM client for response generation
        let response = self.generate_optimized_response().await?;

        // Update conversation history
        self.update_conversation_history(response).await?;

        // Transition back to idle
        self.transition_state(
            AgentState::GeneratingResponse,
            AgentState::Idle,
            TransitionTrigger::ToolCompletion,
        )
        .await?;

        Ok(())
    }

    /// Handle waiting for user state
    async fn handle_waiting_state(&self) -> Result<()> {
        debug!("Waiting for user input");
        tokio::time::sleep(Duration::from_millis(100)).await;
        Ok(())
    }

    /// Handle error state with recovery
    async fn handle_error_state(&self, error_type: &str) -> Result<()> {
        warn!("Handling error state: {}", error_type);

        // Implement error recovery strategies
        match error_type {
            "tool_timeout" => self.recover_from_tool_timeout().await?,
            "memory_limit" => self.recover_from_memory_limit().await?,
            "rate_limit" => self.recover_from_rate_limit().await?,
            _ => {
                error!("Unknown error type: {}", error_type);
                self.transition_state(
                    AgentState::Error {
                        error_type: error_type.to_string(),
                    },
                    AgentState::Idle,
                    TransitionTrigger::InternalError,
                )
                .await?;
            }
        }

        Ok(())
    }

    /// Execute a group of tools in parallel
    async fn execute_tool_group_parallel(&self, tool_group: &[String]) -> Result<()> {
        debug!("Executing tool group in parallel: {:?}", tool_group);

        let mut handles = Vec::new();

        for tool_name in tool_group {
            let request = ToolRequest {
                call: ToolCallRequest {
                    id: uuid::Uuid::new_v4().to_string(),
                    tool_name: tool_name.clone(),
                    args: Value::Object(serde_json::Map::new()),
                    metadata: None,
                },
                priority: ExecutionPriority::Normal,
                timeout: Duration::from_secs(60),
                context: ExecutionContext {
                    session_id: self.context.read().await.session_id.clone(),
                    user_id: None,
                    workspace_path: "/tmp".to_string(),
                    parent_request_id: None,
                },
            };

            let pipeline = Arc::clone(&self.tool_pipeline);
            let handle = tokio::spawn(async move { pipeline.submit_request(request).await });

            handles.push(handle);
        }

        // Wait for all tools in group to complete
        for handle in handles {
            if let Err(e) = handle.await? {
                warn!("Tool execution failed: {}", e);
            }
        }

        Ok(())
    }

    /// Pre-warm caches based on predicted tool usage
    async fn pre_warm_caches(&self, predicted_tools: &[String]) -> Result<()> {
        debug!("Pre-warming caches for tools: {:?}", predicted_tools);

        // This would implement cache warming logic
        // For now, just log the prediction
        for tool_name in predicted_tools {
            debug!("Cache warming candidate: {}", tool_name);
        }

        Ok(())
    }

    /// Generate optimized response using LLM client
    async fn generate_optimized_response(&self) -> Result<String> {
        // This would use the optimized LLM client
        // For now, return a placeholder
        Ok("Optimized response generated".to_string())
    }

    /// Update conversation history with new turn
    async fn update_conversation_history(&self, response: String) -> Result<()> {
        let mut context = self.context.write().await;

        let turn = ConversationTurn {
            id: uuid::Uuid::new_v4().to_string(),
            user_message: "User input".to_string(), // Would be actual user input
            agent_response: Some(response),
            tool_calls: Vec::new(),
            timestamp: Instant::now(),
            execution_time: None,
        };

        context.conversation_history.push(turn);

        // Keep history bounded
        if context.conversation_history.len() > 100 {
            context.conversation_history.remove(0);
        }

        Ok(())
    }

    /// Perform optimizations while agent is idle
    async fn optimize_while_idle(&self) -> Result<()> {
        // Update performance predictions
        self.predictor.update_predictions().await?;

        // Clean up completed tool states
        self.cleanup_completed_tools().await?;

        // Optimize memory usage
        self.optimize_memory_usage().await?;

        Ok(())
    }

    /// Clean up completed tool execution states
    async fn cleanup_completed_tools(&self) -> Result<()> {
        let mut context = self.context.write().await;

        context.active_tools.retain(|_, state| {
            !matches!(
                state.status,
                ToolStatus::Completed | ToolStatus::Failed { .. }
            )
        });

        Ok(())
    }

    /// Optimize memory usage by cleaning up old data
    async fn optimize_memory_usage(&self) -> Result<()> {
        // Clean up old state history
        let mut history = self.state_history.write().await;
        while history.len() > 500 {
            history.pop_front();
        }

        // Return unused memory to pool
        let _pool = global_pool();
        // This would implement memory cleanup logic

        Ok(())
    }

    /// Transition between states with logging
    async fn transition_state(
        &self,
        from: AgentState,
        to: AgentState,
        trigger: TransitionTrigger,
    ) -> Result<()> {
        debug!(
            "State transition: {:?} -> {:?} (trigger: {:?})",
            from, to, trigger
        );

        // Update context
        {
            let mut context = self.context.write().await;
            context.current_state = to.clone();
        }

        // Record transition
        let transition = StateTransition {
            from_state: from,
            to_state: to,
            trigger,
            timestamp: Instant::now(),
        };

        // Send transition event
        if let Err(e) = self.state_tx.send(transition.clone()) {
            warn!("Failed to send state transition event: {}", e);
        }

        // Add to history
        let mut history = self.state_history.write().await;
        history.push_back(transition);

        Ok(())
    }

    /// Start background state monitoring
    async fn start_state_monitor(&self) -> Result<()> {
        let mut rx_guard = self.state_rx.write().await;
        let state_rx = rx_guard
            .take()
            .ok_or_else(|| anyhow!("State monitor already started"))?;
        drop(rx_guard);

        tokio::spawn(async move {
            let mut rx = state_rx;
            while let Some(transition) = rx.recv().await {
                debug!("State monitor: {:?}", transition);
                // This would implement state monitoring logic
            }
        });

        Ok(())
    }

    /// Start performance predictor updates
    async fn start_performance_predictor(&self) -> Result<()> {
        let predictor = Arc::clone(&self.predictor);

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(30));

            loop {
                interval.tick().await;
                if let Err(e) = predictor.update_predictions().await {
                    warn!("Failed to update performance predictions: {}", e);
                }
            }
        });

        Ok(())
    }

    /// Recovery from tool timeout
    async fn recover_from_tool_timeout(&self) -> Result<()> {
        warn!("Recovering from tool timeout");

        // Cancel timed out tools
        let mut context = self.context.write().await;
        for (_, state) in context.active_tools.iter_mut() {
            if matches!(state.status, ToolStatus::Running) {
                state.status = ToolStatus::Cancelled;
            }
        }

        // Transition back to idle
        drop(context);
        self.transition_state(
            AgentState::Error {
                error_type: "tool_timeout".to_string(),
            },
            AgentState::Idle,
            TransitionTrigger::InternalError,
        )
        .await?;

        Ok(())
    }

    /// Recovery from memory limit
    async fn recover_from_memory_limit(&self) -> Result<()> {
        warn!("Recovering from memory limit");

        // Force memory cleanup
        self.optimize_memory_usage().await?;

        // Reduce concurrent tool limit
        let mut context = self.context.write().await;
        context.resource_limits.max_concurrent_tools =
            (context.resource_limits.max_concurrent_tools / 2).max(1);

        drop(context);
        self.transition_state(
            AgentState::Error {
                error_type: "memory_limit".to_string(),
            },
            AgentState::Idle,
            TransitionTrigger::InternalError,
        )
        .await?;

        Ok(())
    }

    /// Recovery from rate limit
    async fn recover_from_rate_limit(&self) -> Result<()> {
        warn!("Recovering from rate limit");

        // Wait before retrying
        tokio::time::sleep(Duration::from_secs(5)).await;

        self.transition_state(
            AgentState::Error {
                error_type: "rate_limit".to_string(),
            },
            AgentState::Idle,
            TransitionTrigger::InternalError,
        )
        .await?;

        Ok(())
    }
}

impl PerformancePredictor {
    pub fn new() -> Self {
        Self {
            execution_patterns: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Predict execution pattern for a given context
    pub async fn predict_execution_pattern(&self, context: &str) -> Option<ExecutionPattern> {
        let patterns = self.execution_patterns.read().await;
        patterns.get(context).cloned()
    }

    /// Update performance predictions based on historical data
    pub async fn update_predictions(&self) -> Result<()> {
        debug!("Updating performance predictions");

        // This would implement machine learning-based prediction updates
        // For now, just log the update

        Ok(())
    }
}