coderlib 0.1.0

//! AI Agent implementation for CoderLib
//!
//! This module contains the main Agent struct that orchestrates AI interactions,
//! tool execution, and response streaming.

use std::collections::HashMap;
use std::sync::Arc;

use dashmap::DashMap;
use tokio::sync::broadcast;
use tokio::task::AbortHandle;
use uuid::Uuid;

use crate::core::{
    CoderLibConfig, CodeRequest, CodeResponse, ToolCall, TokenUsage,
    CoderLibError,
};
use crate::integration::HostIntegration;
use crate::llm::{Provider, ProviderEvent};
use crate::storage::{Storage, Message, MessageRole, MessageContent};
use crate::tools::{Tool, ToolError};

/// Main AI agent that processes requests and manages conversations
pub struct Agent {
    /// Configuration
    config: CoderLibConfig,
    
    /// LLM providers
    providers: HashMap<String, Box<dyn Provider>>,
    
    /// Available tools
    tools: HashMap<String, Box<dyn Tool>>,
    
    /// Storage backend
    storage: Arc<dyn Storage>,
    
    /// Host integration for callbacks
    host_integration: Option<Arc<dyn HostIntegration>>,
    
    /// Active requests (session_id -> abort_handle)
    active_requests: Arc<DashMap<String, AbortHandle>>,
    
    /// Event broadcaster
    event_sender: broadcast::Sender<AgentEvent>,
}

/// Events emitted by the agent during processing
#[derive(Debug, Clone)]
pub enum AgentEvent {
    /// Request processing started
    RequestStarted {
        session_id: String,
        request_id: String,
    },
    
    /// Received a chunk of response from the LLM
    ResponseChunk {
        session_id: String,
        request_id: String,
        content: String,
    },
    
    /// Tool execution requested
    ToolExecutionStarted {
        session_id: String,
        request_id: String,
        tool_name: String,
        parameters: serde_json::Value,
    },
    
    /// Tool execution completed
    ToolExecutionCompleted {
        session_id: String,
        request_id: String,
        tool_name: String,
        result: String,
    },
    
    /// Request completed successfully
    RequestCompleted {
        session_id: String,
        request_id: String,
        response: String,
        token_usage: Option<TokenUsage>,
    },
    
    /// Request failed with error
    RequestFailed {
        session_id: String,
        request_id: String,
        error: String,
    },
    
    /// Request was cancelled
    RequestCancelled {
        session_id: String,
        request_id: String,
    },
}

impl Agent {
    /// Create a new agent with the given configuration
    pub async fn new(
        config: CoderLibConfig,
        storage: Arc<dyn Storage>,
    ) -> Result<Self, CoderLibError> {
        let (event_sender, _) = broadcast::channel(1000);
        
        let mut agent = Self {
            config: config.clone(),
            providers: HashMap::new(),
            tools: HashMap::new(),
            storage,
            host_integration: None,
            active_requests: Arc::new(DashMap::new()),
            event_sender,
        };
        
        // Initialize providers
        agent.initialize_providers().await?;
        
        // Initialize tools
        agent.initialize_tools().await?;
        
        Ok(agent)
    }
    
    /// Set the host integration
    pub fn set_host_integration(&mut self, host: Arc<dyn HostIntegration>) {
        self.host_integration = Some(host);
    }
    
    /// Process a code request and return a stream of responses
    pub async fn process_request(
        &self,
        request: CodeRequest,
    ) -> Result<broadcast::Receiver<CodeResponse>, CoderLibError> {
        let request_id = Uuid::new_v4().to_string();
        let (response_tx, response_rx) = broadcast::channel(100);
        
        // Process the request directly using the main agent
        let session_id = request.session_id.clone();
        let event_sender = self.event_sender.clone();

        // Process the request directly instead of spawning a task
        // This ensures we have access to the real providers
        if let Err(e) = self.process_request_internal(
            request,
            request_id.clone(),
            response_tx.clone(),
            event_sender,
        ).await {
            tracing::error!("Request processing failed: {}", e);

            // Send error response
            let _ = response_tx.send(CodeResponse {
                session_id: session_id.clone(),
                content: format!("Error processing request: {}", e),
                is_complete: true,
                tool_calls: Vec::new(),
                token_usage: None,
            });
        }

        // Create a dummy handle since we're not actually spawning a task
        let handle = tokio::spawn(async {});
        
        // Store the abort handle
        self.active_requests.insert(session_id, handle.abort_handle());
        
        Ok(response_rx)
    }
    
    /// Cancel an ongoing request
    pub async fn cancel_request(&self, session_id: &str) -> Result<(), CoderLibError> {
        if let Some((_, handle)) = self.active_requests.remove(session_id) {
            handle.abort();
            
            // Emit cancellation event
            let _ = self.event_sender.send(AgentEvent::RequestCancelled {
                session_id: session_id.to_string(),
                request_id: "unknown".to_string(), // We don't track request IDs in the map
            });
        }
        
        Ok(())
    }
    
    /// Get available models from all providers
    pub fn available_models(&self) -> Vec<crate::llm::Model> {
        self.providers
            .values()
            .map(|provider| provider.model().clone())
            .collect()
    }
    
    /// Update the agent configuration
    pub async fn update_config(&mut self, config: CoderLibConfig) -> Result<(), CoderLibError> {
        self.config = config;
        
        // Reinitialize providers and tools
        self.providers.clear();
        self.tools.clear();
        
        self.initialize_providers().await?;
        self.initialize_tools().await?;
        
        Ok(())
    }
    
    /// Internal request processing implementation
    async fn process_request_internal(
        &self,
        request: CodeRequest,
        request_id: String,
        response_tx: broadcast::Sender<CodeResponse>,
        event_sender: broadcast::Sender<AgentEvent>,
    ) -> Result<(), CoderLibError> {
        // Emit start event
        let _ = event_sender.send(AgentEvent::RequestStarted {
            session_id: request.session_id.clone(),
            request_id: request_id.clone(),
        });
        
        // Get the appropriate provider and agent config
        let agent_config = self.config.default_agent()?;
        let provider = self.providers.get(&agent_config.provider)
            .ok_or_else(|| CoderLibError::InvalidInput(
                format!("Provider '{}' not found", agent_config.provider)
            ))?;
        
        // Build the conversation history
        let messages = self.build_conversation_history(&request).await?;
        
        // Get available tools for this agent
        let available_tools: Vec<Box<dyn Tool>> = agent_config.tools
            .iter()
            .filter_map(|tool_name| self.tools.get(tool_name).map(|t| t.clone_box()))
            .collect();
        
        // Send request to provider
        if agent_config.streaming {
            self.handle_streaming_response(
                provider.as_ref(),
                messages,
                available_tools,
                &request,
                &request_id,
                response_tx,
                event_sender,
            ).await?;
        } else {
            self.handle_non_streaming_response(
                provider.as_ref(),
                messages,
                available_tools,
                &request,
                &request_id,
                response_tx,
                event_sender,
            ).await?;
        }
        
        Ok(())
    }
    
    /// Handle streaming response from provider
    async fn handle_streaming_response(
        &self,
        provider: &dyn Provider,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
        request: &CodeRequest,
        request_id: &str,
        response_tx: broadcast::Sender<CodeResponse>,
        event_sender: broadcast::Sender<AgentEvent>,
    ) -> Result<(), CoderLibError> {
        let mut stream = provider.stream_response(messages, tools).await?;
        let mut accumulated_content = String::new();
        let mut tool_calls = Vec::new();
        
        while let Some(event) = stream.recv().await {
            match event {
                Ok(ProviderEvent::ContentChunk { content }) => {
                    accumulated_content.push_str(&content);
                    
                    // Emit chunk event
                    let _ = event_sender.send(AgentEvent::ResponseChunk {
                        session_id: request.session_id.clone(),
                        request_id: request_id.to_string(),
                        content: content.clone(),
                    });
                    
                    // Send response chunk
                    let _ = response_tx.send(CodeResponse {
                        session_id: request.session_id.clone(),
                        content,
                        is_complete: false,
                        tool_calls: Vec::new(),
                        token_usage: None,
                    });
                }
                Ok(ProviderEvent::ToolCall { name, parameters }) => {
                    // Execute tool if we have host integration
                    if let Some(host) = &self.host_integration {
                        let tool_result = self.execute_tool(&name, &parameters, host.as_ref()).await?;
                        
                        tool_calls.push(ToolCall {
                            id: Uuid::new_v4().to_string(),
                            name: name.clone(),
                            parameters,
                            result: Some(tool_result),
                        });
                    }
                }
                Ok(ProviderEvent::Complete { token_usage }) => {
                    // Emit completion event
                    let _ = event_sender.send(AgentEvent::RequestCompleted {
                        session_id: request.session_id.clone(),
                        request_id: request_id.to_string(),
                        response: accumulated_content.clone(),
                        token_usage: token_usage.clone(),
                    });

                    // Send final response
                    let _ = response_tx.send(CodeResponse {
                        session_id: request.session_id.clone(),
                        content: accumulated_content,
                        is_complete: true,
                        tool_calls,
                        token_usage,
                    });

                    break;
                }
                Ok(ProviderEvent::Error { error }) => {
                    let _ = event_sender.send(AgentEvent::RequestFailed {
                        session_id: request.session_id.clone(),
                        request_id: request_id.to_string(),
                        error: error.clone(),
                    });

                    return Err(CoderLibError::Internal(error));
                }
                Err(e) => {
                    let error_msg = e.to_string();
                    
                    // Emit error event
                    let _ = event_sender.send(AgentEvent::RequestFailed {
                        session_id: request.session_id.clone(),
                        request_id: request_id.to_string(),
                        error: error_msg.clone(),
                    });
                    
                    return Err(CoderLibError::Provider(e));
                }
            }
        }
        
        Ok(())
    }
    
    /// Handle non-streaming response from provider
    async fn handle_non_streaming_response(
        &self,
        provider: &dyn Provider,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
        request: &CodeRequest,
        request_id: &str,
        response_tx: broadcast::Sender<CodeResponse>,
        event_sender: broadcast::Sender<AgentEvent>,
    ) -> Result<(), CoderLibError> {
        let response = provider.send_messages(messages, tools).await?;
        
        // Process any tool calls
        let mut tool_calls = Vec::new();
        for tool_call in response.tool_calls {
            if let Some(host) = &self.host_integration {
                let tool_result = self.execute_tool(&tool_call.name, &tool_call.parameters, host.as_ref()).await?;
                
                tool_calls.push(ToolCall {
                    id: tool_call.id,
                    name: tool_call.name,
                    parameters: tool_call.parameters,
                    result: Some(tool_result),
                });
            }
        }
        
        // Emit completion event
        let _ = event_sender.send(AgentEvent::RequestCompleted {
            session_id: request.session_id.clone(),
            request_id: request_id.to_string(),
            response: response.content.clone(),
            token_usage: response.token_usage.clone(),
        });
        
        // Send response
        let _ = response_tx.send(CodeResponse {
            session_id: request.session_id.clone(),
            content: response.content,
            is_complete: true,
            tool_calls,
            token_usage: response.token_usage,
        });
        
        Ok(())
    }
    
    /// Execute a tool with the given parameters
    async fn execute_tool(
        &self,
        tool_name: &str,
        parameters: &serde_json::Value,
        host: &dyn HostIntegration,
    ) -> Result<String, CoderLibError> {
        let tool = self.tools.get(tool_name)
            .ok_or_else(|| CoderLibError::InvalidInput(format!("Tool '{}' not found", tool_name)))?;
        
        // Check permissions
        let permission = tool.requires_permission();
        if !host.request_permission(permission).await? {
            return Err(CoderLibError::Tool(ToolError::PermissionDenied(
                format!("Permission denied for tool '{}'", tool_name)
            )));
        }

        // Execute the tool
        let result = tool.execute(parameters.clone(), host).await
            .map_err(|e| CoderLibError::Tool(e))?;
        
        Ok(result.content)
    }
    
    /// Build conversation history for the request
    async fn build_conversation_history(&self, request: &CodeRequest) -> Result<Vec<Message>, CoderLibError> {
        // Get existing messages from storage
        let mut messages = self.storage.get_messages(&request.session_id).await?;
        
        // Add the current request as a user message
        let user_message = Message {
            id: Uuid::new_v4().to_string(),
            session_id: request.session_id.clone(),
            role: MessageRole::User,
            content: MessageContent::Text(request.content.clone()),
            timestamp: chrono::Utc::now(),
            metadata: serde_json::json!({
                "context": request.context,
                "attachments": request.attachments,
            }),
        };
        
        messages.push(user_message);
        
        Ok(messages)
    }
    
    /// Initialize LLM providers based on configuration
    async fn initialize_providers(&mut self) -> Result<(), CoderLibError> {
        use crate::llm::models;

        // Initialize real providers based on configuration
        for (name, provider_config) in &self.config.providers {
            if provider_config.enabled && provider_config.api_key.is_some() {
                let api_key = provider_config.api_key.as_ref().unwrap().clone();

                // Get the appropriate model for this provider
                let model = if let Some(model) = models::find_model_by_id(&provider_config.default_model) {
                    model
                } else {
                    // Fallback to default model for provider
                    models::default_model_for_provider(name)
                        .ok_or_else(|| CoderLibError::Config(
                            crate::core::error::ConfigError::InvalidValue(
                                format!("No default model found for provider '{}'", name)
                            )
                        ))?
                };

                let model_id = model.id.clone(); // Store model ID for logging

                match name.as_str() {
                    "openai" => {
                        let provider = crate::llm::providers::OpenAIProvider::new(
                            api_key,
                            model,
                            provider_config.base_url.clone(),
                            provider_config.max_tokens,
                            Some(0.1), // Default temperature
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized OpenAI provider with model: {}", model_id);
                    }
                    "anthropic" => {
                        let provider = crate::llm::providers::AnthropicProvider::new(
                            api_key,
                            model,
                            provider_config.base_url.clone(),
                            provider_config.max_tokens,
                            Some(0.1), // Default temperature
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized Anthropic provider with model: {}", model_id);
                    }
                    "gemini" => {
                        let provider = crate::llm::providers::GeminiProvider::new(
                            api_key,
                            model,
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized Gemini provider with model: {}", model_id);
                    }
                    "groq" => {
                        let provider = crate::llm::providers::OpenAICompatibleProvider::groq(
                            api_key,
                            model,
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized Groq provider with model: {}", model_id);
                    }
                    "cohere" => {
                        let provider = crate::llm::providers::OpenAICompatibleProvider::cohere(
                            api_key,
                            model,
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized Cohere provider with model: {}", model_id);
                    }
                    "sambanova" => {
                        let provider = crate::llm::providers::OpenAICompatibleProvider::sambanova(
                            api_key,
                            model,
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized SambaNova provider with model: {}", model_id);
                    }
                    "together" => {
                        let provider = crate::llm::providers::OpenAICompatibleProvider::together(
                            api_key,
                            model,
                        )?;
                        self.providers.insert(name.clone(), Box::new(provider));
                        eprintln!("✅ Initialized Together provider with model: {}", model_id);
                    }
                    _ => {
                        eprintln!("⚠️ Unknown provider '{}', skipping", name);
                        continue;
                    }
                }
            }
        }

        // If no providers were initialized, add a mock provider as fallback
        if self.providers.is_empty() {
            eprintln!("⚠️ No providers configured, using mock provider");
            let mock_provider = crate::llm::providers::MockProvider::new();
            self.providers.insert("mock".to_string(), Box::new(mock_provider));
        } else {
            eprintln!("🚀 Successfully initialized {} real LLM provider(s)", self.providers.len());
        }

        Ok(())
    }

    /// Initialize tools based on configuration
    async fn initialize_tools(&mut self) -> Result<(), CoderLibError> {
        use crate::tools::*;

        // Initialize file operation tools
        if self.config.tools.file_operations_enabled {
            self.tools.insert("file_read".to_string(), Box::new(file_ops::FileReadTool::new()));
            self.tools.insert("file_write".to_string(), Box::new(file_ops::FileWriteTool::new()));
            self.tools.insert("file_list".to_string(), Box::new(file_ops::FileListTool::new()));
        }

        // Initialize shell tools
        if self.config.tools.shell_enabled {
            self.tools.insert("shell_command".to_string(), Box::new(shell::ShellCommandTool::new()));
        }

        // Initialize search tools
        self.tools.insert("file_search".to_string(), Box::new(search::FileSearchTool::new()));
        self.tools.insert("content_search".to_string(), Box::new(search::ContentSearchTool::new()));

        Ok(())
    }
    
}

#[cfg(test)]
mod tests {
    use super::*;
    // use crate::core::CoderLibConfig;

    #[test]
    fn test_agent_event_creation() {
        let event = AgentEvent::RequestStarted {
            session_id: "test-session".to_string(),
            request_id: "test-request".to_string(),
        };
        
        match event {
            AgentEvent::RequestStarted { session_id, request_id } => {
                assert_eq!(session_id, "test-session");
                assert_eq!(request_id, "test-request");
            }
            _ => panic!("Expected RequestStarted event"),
        }
    }
}