QuiCommit/src/llm/openai.rs

use super::{create_http_client, LlmProvider};
use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// OpenAI API client
pub struct OpenAiClient {
    base_url: String,
    api_key: String,
    model: String,
    client: reqwest::Client,
}

#[derive(Debug, Serialize)]
struct ChatCompletionRequest {
    model: String,
    messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    stream: bool,
}

#[derive(Debug, Serialize, Deserialize)]
struct Message {
    role: String,
    content: String,
}

#[derive(Debug, Deserialize)]
struct ChatCompletionResponse {
    choices: Vec<Choice>,
}

#[derive(Debug, Deserialize)]
struct Choice {
    message: Message,
}

#[derive(Debug, Deserialize)]
struct ErrorResponse {
    error: ApiError,
}

#[derive(Debug, Deserialize)]
struct ApiError {
    message: String,
    #[serde(rename = "type")]
    error_type: String,
}

impl OpenAiClient {
    /// Create new OpenAI client
    pub fn new(base_url: &str, api_key: &str, model: &str) -> Result<Self> {
        let client = create_http_client(Duration::from_secs(60))?;
        
        Ok(Self {
            base_url: base_url.trim_end_matches('/').to_string(),
            api_key: api_key.to_string(),
            model: model.to_string(),
            client,
        })
    }

    /// Set timeout
    pub fn with_timeout(mut self, timeout: Duration) -> Result<Self> {
        self.client = create_http_client(timeout)?;
        Ok(self)
    }

    /// List available models
    pub async fn list_models(&self) -> Result<Vec<String>> {
        let url = format!("{}/models", self.base_url);
        
        let response = self.client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .send()
            .await
            .context("Failed to list OpenAI models")?;
        
        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.unwrap_or_default();
            bail!("OpenAI API error: {} - {}", status, text);
        }
        
        #[derive(Deserialize)]
        struct ModelsResponse {
            data: Vec<Model>,
        }
        
        #[derive(Deserialize)]
        struct Model {
            id: String,
        }
        
        let result: ModelsResponse = response
            .json()
            .await
            .context("Failed to parse OpenAI response")?;
        
        Ok(result.data.into_iter().map(|m| m.id).collect())
    }

    /// Validate API key
    pub async fn validate_key(&self) -> Result<bool> {
        match self.list_models().await {
            Ok(_) => Ok(true),
            Err(e) => {
                let err_str = e.to_string();
                if err_str.contains("401") || err_str.contains("Unauthorized") {
                    Ok(false)
                } else {
                    Err(e)
                }
            }
        }
    }
}

#[async_trait]
impl LlmProvider for OpenAiClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let messages = vec![
            Message {
                role: "user".to_string(),
                content: prompt.to_string(),
            },
        ];
        
        self.chat_completion(messages).await
    }

    async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
        let mut messages = vec![];
        
        if !system.is_empty() {
            messages.push(Message {
                role: "system".to_string(),
                content: system.to_string(),
            });
        }
        
        messages.push(Message {
            role: "user".to_string(),
            content: user.to_string(),
        });
        
        self.chat_completion(messages).await
    }

    async fn is_available(&self) -> bool {
        self.validate_key().await.unwrap_or(false)
    }

    fn name(&self) -> &str {
        "openai"
    }
}

impl OpenAiClient {
    async fn chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        let url = format!("{}/chat/completions", self.base_url);
        
        let request = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            max_tokens: Some(500),
            temperature: Some(0.7),
            stream: false,
        };
        
        let response = self.client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send request to OpenAI")?;
        
        let status = response.status();
        
        if !status.is_success() {
            let text = response.text().await.unwrap_or_default();
            
            // Try to parse error
            if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
                bail!("OpenAI API error: {} ({})", error.error.message, error.error.error_type);
            }
            
            bail!("OpenAI API error: {} - {}", status, text);
        }
        
        let result: ChatCompletionResponse = response
            .json()
            .await
            .context("Failed to parse OpenAI response")?;
        
        result.choices
            .into_iter()
            .next()
            .map(|c| c.message.content.trim().to_string())
            .ok_or_else(|| anyhow::anyhow!("No response from OpenAI"))
    }
}

/// Azure OpenAI client (extends OpenAI with Azure-specific config)
pub struct AzureOpenAiClient {
    endpoint: String,
    api_key: String,
    deployment: String,
    api_version: String,
    client: reqwest::Client,
}

impl AzureOpenAiClient {
    /// Create new Azure OpenAI client
    pub fn new(
        endpoint: &str,
        api_key: &str,
        deployment: &str,
        api_version: &str,
    ) -> Result<Self> {
        let client = create_http_client(Duration::from_secs(60))?;
        
        Ok(Self {
            endpoint: endpoint.trim_end_matches('/').to_string(),
            api_key: api_key.to_string(),
            deployment: deployment.to_string(),
            api_version: api_version.to_string(),
            client,
        })
    }

    async fn chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        let url = format!(
            "{}/openai/deployments/{}/chat/completions?api-version={}",
            self.endpoint, self.deployment, self.api_version
        );
        
        let request = ChatCompletionRequest {
            model: self.deployment.clone(),
            messages,
            max_tokens: Some(500),
            temperature: Some(0.7),
            stream: false,
        };
        
        let response = self.client
            .post(&url)
            .header("api-key", &self.api_key)
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send request to Azure OpenAI")?;
        
        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.unwrap_or_default();
            bail!("Azure OpenAI API error: {} - {}", status, text);
        }
        
        let result: ChatCompletionResponse = response
            .json()
            .await
            .context("Failed to parse Azure OpenAI response")?;
        
        result.choices
            .into_iter()
            .next()
            .map(|c| c.message.content.trim().to_string())
            .ok_or_else(|| anyhow::anyhow!("No response from Azure OpenAI"))
    }
}

#[async_trait]
impl LlmProvider for AzureOpenAiClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let messages = vec![
            Message {
                role: "user".to_string(),
                content: prompt.to_string(),
            },
        ];
        
        self.chat_completion(messages).await
    }

    async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
        let mut messages = vec![];
        
        if !system.is_empty() {
            messages.push(Message {
                role: "system".to_string(),
                content: system.to_string(),
            });
        }
        
        messages.push(Message {
            role: "user".to_string(),
            content: user.to_string(),
        });
        
        self.chat_completion(messages).await
    }

    async fn is_available(&self) -> bool {
        // Simple check - try to make a minimal request
        let url = format!(
            "{}/openai/deployments/{}/chat/completions?api-version={}",
            self.endpoint, self.deployment, self.api_version
        );
        
        let request = ChatCompletionRequest {
            model: self.deployment.clone(),
            messages: vec![Message {
                role: "user".to_string(),
                content: "Hi".to_string(),
            }],
            max_tokens: Some(5),
            temperature: Some(0.0),
            stream: false,
        };
        
        match self.client
            .post(&url)
            .header("api-key", &self.api_key)
            .json(&request)
            .send()
            .await
        {
            Ok(response) => response.status().is_success(),
            Err(_) => false,
        }
    }

    fn name(&self) -> &str {
        "azure-openai"
    }
}