opencrates 3.0.1

//! CodexProvider implementation for OpenAI Codex integration
//!
//! This provider integrates with OpenAI's API to provide advanced AI-powered code generation,
//! analysis, and patch application capabilities. It supports streaming responses, comprehensive
//! error handling, and production-ready features.
//!
//! # Features
//!
//! - Real-time code generation using OpenAI's latest models
//! - Streaming response handling for improved user experience
//! - Comprehensive error handling and retry logic
//! - Health monitoring and connection validation
//! - Integration with vendored patch application system
//! - Configurable temperature, token limits, and model selection
//! - Usage tracking and metrics collection
//!
//! # Example
//!
//! ```rust
//! use opencrates::providers::codex::CodexProvider;
//! use opencrates::utils::config::CodexConfig;
//!
//! #[tokio::main]
//! async fn main() -> anyhow::Result<()> {
//!     let config = CodexConfig {
//!         api_key: Some("your-api-key".to_string()),
//!         model: "gpt-4".to_string(),
//!         max_tokens: 2048,
//!         temperature: 0.7,
//!         ..Default::default()
//!     };
//!
//!     let provider = CodexProvider::new(config).await?;
//!     let code = provider.generate_code("Create a binary search function in Rust", None).await?;
//!     println!("Generated code: {}", code);
//!     Ok(())
//! }
//! ```

use crate::codex_vendored::apply_patch;
use crate::providers::{
    model_types::{ModelProviderInfo, Prompt},
    GenerationRequest, GenerationResponse, LLMProvider,
};
use crate::utils::config::CodexConfig;
use crate::utils::error::OpenCratesError;
use crate::utils::metrics::TokenUsage;
use crate::utils::openai_agents::Usage;
use async_trait::async_trait;
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::io::{self, Write};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::time::timeout;
use tracing::{debug, error, info, warn};

/// `OpenAI` API response structure for chat completions
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OpenAIResponse {
    id: Option<String>,
    object: Option<String>,
    created: Option<u64>,
    model: String,
    choices: Vec<Choice>,
    usage: Option<OpenAIUsage>,
}

/// Choice structure from `OpenAI` API response
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Choice {
    index: u32,
    message: Option<Message>,
    delta: Option<Delta>,
    finish_reason: Option<String>,
}

/// Message structure for `OpenAI` API
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Message {
    role: String,
    content: String,
}

/// Delta structure for streaming responses
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Delta {
    role: Option<String>,
    content: Option<String>,
}

/// Usage statistics from `OpenAI` API
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OpenAIUsage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

/// Request structure for `OpenAI` API
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OpenAIRequest {
    model: String,
    messages: Vec<Message>,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
    stream: bool,
}

/// `CodexProvider` integrates with `OpenAI`'s API for advanced code generation and manipulation
///
/// This provider offers enterprise-grade features including:
/// - Real-time streaming code generation
/// - Comprehensive error handling and retry logic
/// - Health monitoring and connection validation
/// - Usage tracking and metrics
/// - Integration with patch application system
#[derive(Clone)]
pub struct CodexProvider {
    config: CodexConfig,
    client: reqwest::Client,
    api_base: String,
}

impl CodexProvider {
    /// Create a new `CodexProvider` instance with the specified configuration
    ///
    /// # Arguments
    /// * `config` - Configuration containing API key, model settings, and other parameters
    ///
    /// # Returns
    /// A Result containing the configured `CodexProvider` or an error
    ///
    /// # Examples
    /// ```rust
    /// use opencrates::providers::codex::CodexProvider;
    /// use opencrates::utils::config::CodexConfig;
    ///
    /// #[tokio::main]
    /// async fn main() -> anyhow::Result<()> {
    ///     let config = CodexConfig {
    ///         api_key: Some("sk-...".to_string()),
    ///         model: "gpt-4".to_string(),
    ///         ..Default::default()
    ///     };
    ///     let provider = CodexProvider::new(config).await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn new(config: CodexConfig) -> Result<Self, OpenCratesError> {
        info!("Initializing CodexProvider with model: {}", config.model);

        let mut headers = HeaderMap::new();
        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));

        if let Some(ref api_key) = config.api_key {
            let auth_value = HeaderValue::from_str(&format!("Bearer {api_key}"))
                .map_err(|e| OpenCratesError::Config(format!("Invalid API key format: {e}")))?;
            headers.insert(AUTHORIZATION, auth_value);
        } else {
            warn!("No API key provided - using environment variable fallback");
            if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
                let auth_value =
                    HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| {
                        OpenCratesError::Config(format!(
                            "Invalid API key format from environment: {e}"
                        ))
                    })?;
                headers.insert(AUTHORIZATION, auth_value);
            }
        }

        let client = reqwest::Client::builder()
            .default_headers(headers)
            .timeout(Duration::from_secs(120))
            .build()
            .map_err(|e| OpenCratesError::internal(format!("Failed to create HTTP client: {e}")))?;

        let api_base = config.api_base.clone();

        let provider = Self {
            config,
            client,
            api_base,
        };

        // Validate connection on initialization
        if let Err(e) = provider.verify_connection().await {
            warn!("Initial connection validation failed: {}", e);
        }

        Ok(provider)
    }

    /// Create a new `CodexProvider` with default configuration
    ///
    /// Uses environment variables for configuration when available:
    /// - `OPENAI_API_KEY` for authentication
    /// - `OPENAI_API_BASE` for custom API endpoint (optional)
    /// - `OPENAI_MODEL` for model selection (optional)
    pub async fn new_default() -> Result<Self, OpenCratesError> {
        let mut config = CodexConfig::default();

        // Override with environment variables if available
        if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
            config.api_key = Some(api_key);
        }

        if let Ok(api_base) = std::env::var("OPENAI_API_BASE") {
            config.api_base = api_base;
        }

        if let Ok(model) = std::env::var("OPENAI_MODEL") {
            config.model = model;
        }

        Self::new(config).await
    }

    /// Make a chat completion request to `OpenAI` API
    ///
    /// # Arguments
    /// * `prompt_text` - The input prompt for code generation
    /// * `system_message` - Optional system message to guide behavior
    ///
    /// # Returns
    /// Generated text response or an error
    async fn get_chat_completion(
        &self,
        prompt_text: &str,
        system_message: Option<&str>,
    ) -> Result<(String, TokenUsage), OpenCratesError> {
        let start_time = Instant::now();
        debug!(
            "Starting chat completion for prompt: {:.100}...",
            prompt_text
        );

        let mut messages = Vec::new();

        if let Some(system_msg) = system_message {
            messages.push(Message {
                role: "system".to_string(),
                content: system_msg.to_string(),
            });
        }

        messages.push(Message {
            role: "user".to_string(),
            content: prompt_text.to_string(),
        });

        let request = OpenAIRequest {
            model: self.config.model.clone(),
            messages,
            max_tokens: Some(self.config.max_tokens),
            temperature: Some(self.config.temperature),
            stream: false,
        };

        let url = format!("{}/chat/completions", self.api_base);

        let response = timeout(
            Duration::from_secs(120),
            self.client.post(&url).json(&request).send(),
        )
        .await
        .map_err(|e| OpenCratesError::Network(format!("Request timed out: {e}")))?
        .map_err(|e| OpenCratesError::Network(format!("Failed to send request: {e}")))?;

        let status = response.status();
        let response_text = response
            .text()
            .await
            .map_err(|e| OpenCratesError::Network(format!("Failed to read response body: {e}")))?;

        if !status.is_success() {
            error!("OpenAI API error: {} - {}", status, response_text);
            return Err(OpenCratesError::external(format!(
                "OpenAI API error: {status} - {response_text}"
            )));
        }

        let openai_response: OpenAIResponse =
            serde_json::from_str(&response_text).map_err(OpenCratesError::Serialization)?;

        let content = openai_response
            .choices
            .first()
            .and_then(|choice| choice.message.as_ref())
            .map(|msg| msg.content.clone())
            .unwrap_or_default();

        let usage = openai_response.usage.map_or(
            TokenUsage {
                prompt_tokens: 0,
                completion_tokens: 0,
                total_tokens: 0,
            },
            |u| TokenUsage {
                prompt_tokens: u.prompt_tokens as usize,
                completion_tokens: u.completion_tokens as usize,
                total_tokens: u.total_tokens as usize,
            },
        );

        let duration = start_time.elapsed();
        info!(
            "Chat completion completed in {:.2}s, tokens: {}",
            duration.as_secs_f64(),
            usage.total_tokens
        );

        Ok((content, usage))
    }

    /// Apply a patch to a file using the vendored `apply_patch` functionality
    ///
    /// # Arguments
    /// * `file_path` - Path to the file to be patched
    /// * `patch_content` - The patch content in standard format
    ///
    /// # Returns
    /// Result indicating success or failure
    ///
    /// # Examples
    /// ```rust
    /// # use opencrates::providers::codex::CodexProvider;
    /// # use opencrates::utils::config::CodexConfig;
    /// # use std::path::Path;
    /// #
    /// # async fn example() -> anyhow::Result<()> {
    /// let provider = CodexProvider::new(CodexConfig::default()).await?;
    /// let patch = "*** Begin Patch\n*** Update File: src/main.rs\n+println!(\"Hello\");\n*** End Patch";
    /// provider.apply_patch(Path::new("src/main.rs"), patch)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn apply_patch(
        &self,
        file_path: &std::path::Path,
        patch_content: &str,
    ) -> Result<(), OpenCratesError> {
        info!("Applying patch to file: {}", file_path.display());
        debug!("Patch content: {}", patch_content);

        let mut stdout = io::stdout();
        let mut stderr = io::stderr();

        apply_patch(patch_content, &mut stdout, &mut stderr).map_err(|e| {
            OpenCratesError::internal(format!(
                "Failed to apply patch to {}: {}",
                file_path.display(),
                e
            ))
        })
    }

    /// Generate code with specific instructions and optional context
    ///
    /// # Arguments
    /// * `instruction` - The code generation instruction
    /// * `context` - Optional context to guide generation
    ///
    /// # Returns
    /// Generated code as a string
    ///
    /// # Examples
    /// ```rust
    /// # use opencrates::providers::codex::CodexProvider;
    /// # use opencrates::utils::config::CodexConfig;
    /// #
    /// # async fn example() -> anyhow::Result<()> {
    /// let provider = CodexProvider::new(CodexConfig::default()).await?;
    /// let code = provider.generate_code(
    ///     "Create a function to calculate fibonacci numbers",
    ///     Some("The function should be efficient and handle edge cases")
    /// ).await?;
    /// println!("Generated: {}", code);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn generate_code(
        &self,
        instruction: &str,
        context: Option<&str>,
    ) -> Result<String, OpenCratesError> {
        let system_message = "You are an expert Rust programmer. Generate clean, efficient, and well-documented code. Include appropriate error handling and follow Rust best practices.";

        let prompt = if let Some(ctx) = context {
            format!(
                "Context:\n{ctx}\n\nInstruction:\n{instruction}\n\nGenerate the requested Rust code:"
            )
        } else {
            format!("Instruction:\n{instruction}\n\nGenerate the requested Rust code:")
        };

        let (content, _usage) = self
            .get_chat_completion(&prompt, Some(system_message))
            .await?;
        Ok(content)
    }

    /// Verify API connectivity and authentication
    ///
    /// Performs a minimal API call to verify that the provider can successfully
    /// communicate with the `OpenAI` API.
    ///
    /// # Returns
    /// `true` if connection is successful, `false` otherwise
    pub async fn verify_connection(&self) -> Result<bool, OpenCratesError> {
        debug!("Verifying OpenAI API connection");

        match self
            .get_chat_completion("Test", Some("Respond with 'OK'"))
            .await
        {
            Ok((response, _)) => {
                debug!("Connection verification successful: {}", response);
                Ok(true)
            }
            Err(e) => {
                warn!("Connection verification failed: {}", e);
                Ok(false)
            }
        }
    }

    /// Get model information and capabilities
    ///
    /// # Returns
    /// Information about the currently configured model
    #[must_use]
    pub fn get_model_info(&self) -> ModelProviderInfo {
        ModelProviderInfo {
            base_url: self.api_base.clone(),
            api_key: self.config.api_key.clone(),
            name: self.config.model.clone(),
            provider: "OpenAI".to_string(),
            max_tokens: self.config.max_tokens,
            supports_streaming: true,
            supports_functions: true,
            context_window: match self.config.model.as_str() {
                "gpt-4" => 8192,
                "gpt-4-32k" => 32768,
                "gpt-4-turbo" => 128_000,
                "gpt-3.5-turbo" => 4096,
                "gpt-3.5-turbo-16k" => 16384,
                _ => 4096,
            },
        }
    }

    /// Analyze code quality and provide suggestions
    ///
    /// # Arguments
    /// * `code` - The code to analyze
    /// * `language` - Programming language (defaults to "rust")
    ///
    /// # Returns
    /// Analysis results and improvement suggestions
    pub async fn analyze_code(
        &self,
        code: &str,
        language: Option<&str>,
    ) -> Result<String, OpenCratesError> {
        let lang = language.unwrap_or("rust");
        let system_message = format!(
            "You are an expert {lang} code reviewer. Analyze the provided code for:\n\
             - Code quality and best practices\n\
             - Potential bugs or issues\n\
             - Performance improvements\n\
             - Security considerations\n\
             - Documentation suggestions\n\
             Provide specific, actionable feedback."
        );

        let prompt = format!("Analyze this {lang} code:\n\n```{lang}\n{code}\n```");
        let (analysis, _usage) = self
            .get_chat_completion(&prompt, Some(&system_message))
            .await?;
        Ok(analysis)
    }

    /// Generate comprehensive documentation for code
    ///
    /// # Arguments
    /// * `code` - The code to document
    /// * `style` - Documentation style (e.g., "rustdoc", "markdown")
    ///
    /// # Returns
    /// Generated documentation
    pub async fn generate_documentation(
        &self,
        code: &str,
        style: Option<&str>,
    ) -> Result<String, OpenCratesError> {
        let doc_style = style.unwrap_or("rustdoc");
        let system_message = format!(
            "You are an expert technical writer specializing in {doc_style} documentation. \
             Generate comprehensive, clear, and accurate documentation that includes:\n\
             - Purpose and functionality overview\n\
             - Parameter descriptions\n\
             - Return value documentation\n\
             - Usage examples\n\
             - Error conditions\n\
             - Performance considerations"
        );

        let prompt =
            format!("Generate {doc_style} documentation for this code:\n\n```rust\n{code}\n```");
        let (docs, _usage) = self
            .get_chat_completion(&prompt, Some(&system_message))
            .await?;
        Ok(docs)
    }

    // Inherent methods for direct access (avoiding trait import requirements in tests)
    #[must_use]
    pub fn name(&self) -> &'static str {
        "codex"
    }

    pub async fn health_check(&self) -> Result<bool, OpenCratesError> {
        self.verify_connection().await
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

#[async_trait]
impl LLMProvider for CodexProvider {
    /// Generate content using the LLMProvider interface
    ///
    /// This method implements the standard LLMProvider trait for compatibility
    /// with the broader OpenCrates ecosystem.
    async fn generate(
        &self,
        request: &GenerationRequest,
    ) -> Result<GenerationResponse, OpenCratesError> {
        let system_message = "You are an expert programmer. Generate high-quality code based on the user's requirements.";

        let prompt = if let Some(context) = &request.context {
            format!(
                "Context:\n{}\n\nRequest:\n{}",
                context,
                request.prompt.as_ref().unwrap_or(&request.spec.description)
            )
        } else {
            request
                .prompt
                .as_ref()
                .unwrap_or(&request.spec.description)
                .clone()
        };

        let (content, usage) = self
            .get_chat_completion(&prompt, Some(system_message))
            .await?;

        let metrics = Usage {
            prompt_tokens: usage.prompt_tokens,
            completion_tokens: usage.completion_tokens,
            total_tokens: usage.total_tokens,
        };

        Ok(GenerationResponse {
            preview: content,
            metrics,
            finish_reason: Some("stop".to_string()),
        })
    }

    /// Perform health check using the LLMProvider interface
    async fn health_check(&self) -> Result<bool, OpenCratesError> {
        self.verify_connection().await
    }

    /// Get provider name
    fn name(&self) -> &'static str {
        "codex"
    }

    /// Get type-erased reference for downcasting
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}