coderlib 0.1.0

//! OpenAI-compatible provider implementation for CoderLib
//!
//! This module provides a generic provider that works with any service
//! that implements the OpenAI Chat Completions API format, including:
//! - Groq
//! - Cohere (via OpenAI compatibility)
//! - SambaNova
//! - Together AI
//! - Perplexity
//! - And many others

use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use tokio_stream::StreamExt;

use crate::core::error::ProviderError;
use crate::core::TokenUsage;
use crate::llm::{Provider, ProviderEvent, ProviderResponse, ProviderToolCall, Model, ModelCapabilities};
use crate::storage::Message;
use crate::tools::Tool;

/// OpenAI-compatible provider for various services
#[derive(Debug, Clone)]
pub struct OpenAICompatibleProvider {
    client: Client,
    api_key: String,
    base_url: String,
    model: Model,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
    provider_name: String,
    supports_tools: bool,
    supports_streaming: bool,
}

/// Provider configuration for OpenAI-compatible services
#[derive(Debug, Clone)]
pub struct CompatibleProviderConfig {
    pub name: String,
    pub base_url: String,
    pub supports_tools: bool,
    pub supports_streaming: bool,
    pub requires_auth: bool,
    pub custom_headers: Vec<(String, String)>,
}

impl CompatibleProviderConfig {
    /// Groq configuration
    pub fn groq() -> Self {
        Self {
            name: "groq".to_string(),
            base_url: "https://api.groq.com/openai/v1".to_string(),
            supports_tools: true,
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }

    /// Cohere configuration (via OpenAI compatibility)
    pub fn cohere() -> Self {
        Self {
            name: "cohere".to_string(),
            base_url: "https://api.cohere.ai/v1".to_string(),
            supports_tools: true,
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }

    /// SambaNova configuration
    pub fn sambanova() -> Self {
        Self {
            name: "sambanova".to_string(),
            base_url: "https://api.sambanova.ai/v1".to_string(),
            supports_tools: false, // Check if they support tools
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }

    /// Together AI configuration
    pub fn together() -> Self {
        Self {
            name: "together".to_string(),
            base_url: "https://api.together.xyz/v1".to_string(),
            supports_tools: true,
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }

    /// Perplexity configuration
    pub fn perplexity() -> Self {
        Self {
            name: "perplexity".to_string(),
            base_url: "https://api.perplexity.ai".to_string(),
            supports_tools: false,
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }

    /// Custom provider configuration
    pub fn custom(name: String, base_url: String) -> Self {
        Self {
            name,
            base_url,
            supports_tools: true,
            supports_streaming: true,
            requires_auth: true,
            custom_headers: vec![],
        }
    }
}

/// OpenAI chat completion request (reused from openai.rs structure)
#[derive(Debug, Serialize)]
struct ChatCompletionRequest {
    model: String,
    messages: Vec<OpenAIMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<OpenAITool>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_choice: Option<String>,
}

/// OpenAI message format
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIMessage {
    role: String,
    content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OpenAIToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
}

/// OpenAI tool definition
#[derive(Debug, Serialize)]
struct OpenAITool {
    #[serde(rename = "type")]
    tool_type: String,
    function: OpenAIFunction,
}

/// OpenAI function definition
#[derive(Debug, Serialize)]
struct OpenAIFunction {
    name: String,
    description: String,
    parameters: serde_json::Value,
}

/// OpenAI tool call
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIToolCall {
    id: String,
    #[serde(rename = "type")]
    call_type: String,
    function: OpenAIFunctionCall,
}

/// OpenAI function call
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIFunctionCall {
    name: String,
    arguments: String,
}

/// OpenAI chat completion response
#[derive(Debug, Serialize, Deserialize)]
struct ChatCompletionResponse {
    id: String,
    object: String,
    created: u64,
    model: String,
    choices: Vec<Choice>,
    usage: Option<Usage>,
}

/// OpenAI choice
#[derive(Debug, Serialize, Deserialize)]
struct Choice {
    index: u32,
    message: OpenAIMessage,
    finish_reason: Option<String>,
}

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

/// OpenAI streaming chunk
#[derive(Debug, Deserialize)]
struct StreamChunk {
    id: String,
    object: String,
    created: u64,
    model: String,
    choices: Vec<StreamChoice>,
}

/// OpenAI streaming choice
#[derive(Debug, Deserialize)]
struct StreamChoice {
    index: u32,
    delta: Delta,
    finish_reason: Option<String>,
}

/// OpenAI streaming delta
#[derive(Debug, Deserialize)]
struct Delta {
    role: Option<String>,
    content: Option<String>,
    tool_calls: Option<Vec<OpenAIToolCall>>,
}

impl OpenAICompatibleProvider {
    /// Create a new OpenAI-compatible provider
    pub fn new(
        config: CompatibleProviderConfig,
        api_key: String,
        model: Model,
        max_tokens: Option<u32>,
        temperature: Option<f32>,
    ) -> Result<Self, ProviderError> {
        if config.requires_auth && api_key.is_empty() {
            return Err(ProviderError::Configuration("API key is required".into()));
        }

        let mut client_builder = Client::builder()
            .timeout(std::time::Duration::from_secs(60));

        // Add custom headers if specified
        let mut default_headers = reqwest::header::HeaderMap::new();
        for (key, value) in &config.custom_headers {
            if let (Ok(header_name), Ok(header_value)) = (
                reqwest::header::HeaderName::from_bytes(key.as_bytes()),
                reqwest::header::HeaderValue::from_str(value),
            ) {
                default_headers.insert(header_name, header_value);
            }
        }
        
        if !default_headers.is_empty() {
            client_builder = client_builder.default_headers(default_headers);
        }

        let client = client_builder
            .build()
            .map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;

        Ok(Self {
            client,
            api_key,
            base_url: config.base_url,
            model,
            max_tokens,
            temperature,
            provider_name: config.name,
            supports_tools: config.supports_tools,
            supports_streaming: config.supports_streaming,
        })
    }

    /// Create Groq provider
    pub fn groq(api_key: String, model: Model) -> Result<Self, ProviderError> {
        Self::new(CompatibleProviderConfig::groq(), api_key, model, None, None)
    }

    /// Create Cohere provider
    pub fn cohere(api_key: String, model: Model) -> Result<Self, ProviderError> {
        Self::new(CompatibleProviderConfig::cohere(), api_key, model, None, None)
    }

    /// Create SambaNova provider
    pub fn sambanova(api_key: String, model: Model) -> Result<Self, ProviderError> {
        Self::new(CompatibleProviderConfig::sambanova(), api_key, model, None, None)
    }

    /// Create Together AI provider
    pub fn together(api_key: String, model: Model) -> Result<Self, ProviderError> {
        Self::new(CompatibleProviderConfig::together(), api_key, model, None, None)
    }

    /// Create Perplexity provider
    pub fn perplexity(api_key: String, model: Model) -> Result<Self, ProviderError> {
        Self::new(CompatibleProviderConfig::perplexity(), api_key, model, None, None)
    }

    /// Convert CoderLib messages to OpenAI format (reuse logic from openai.rs)
    fn convert_messages(&self, messages: Vec<Message>) -> Vec<OpenAIMessage> {
        messages
            .into_iter()
            .map(|msg| OpenAIMessage {
                role: match msg.role {
                    crate::storage::MessageRole::User => "user".to_string(),
                    crate::storage::MessageRole::Assistant => "assistant".to_string(),
                    crate::storage::MessageRole::System => "system".to_string(),
                    crate::storage::MessageRole::Tool => "tool".to_string(),
                },
                content: match msg.content {
                    crate::storage::MessageContent::Text(text) => text,
                    crate::storage::MessageContent::Structured(parts) => {
                        parts.iter().map(|part| match part {
                            crate::storage::ContentPart::Text { text } => text.clone(),
                            crate::storage::ContentPart::Code { code, language } => {
                                format!("```{}\n{}\n```", language, code)
                            }
                            crate::storage::ContentPart::ToolCall { name, parameters } => {
                                format!("Tool call: {} with parameters: {}", name, parameters)
                            }
                            crate::storage::ContentPart::ToolResult { result, .. } => result.clone(),
                            _ => "[Unsupported content]".to_string(),
                        }).collect::<Vec<_>>().join("\n")
                    }
                    crate::storage::MessageContent::ToolCall { name, parameters, .. } => {
                        format!("Tool call: {} with parameters: {}", name, parameters)
                    }
                    crate::storage::MessageContent::ToolResult { result, .. } => result,
                },
                tool_calls: None,
                tool_call_id: None,
            })
            .collect()
    }

    /// Convert tools to OpenAI format
    fn convert_tools(&self, tools: Vec<Box<dyn Tool>>) -> Vec<OpenAITool> {
        tools
            .into_iter()
            .map(|tool| OpenAITool {
                tool_type: "function".to_string(),
                function: OpenAIFunction {
                    name: tool.name().to_string(),
                    description: tool.description().to_string(),
                    parameters: tool.parameter_schema(),
                },
            })
            .collect()
    }

    /// Convert OpenAI tool calls to CoderLib format
    fn convert_tool_calls(&self, tool_calls: &[OpenAIToolCall]) -> Vec<ProviderToolCall> {
        tool_calls
            .iter()
            .map(|call| ProviderToolCall {
                id: call.id.clone(),
                name: call.function.name.clone(),
                parameters: serde_json::from_str(&call.function.arguments)
                    .unwrap_or(serde_json::Value::Null),
            })
            .collect()
    }

    /// Parse streaming response line (reuse logic from openai.rs)
    fn parse_stream_line(&self, line: &str) -> Option<Result<ProviderEvent, ProviderError>> {
        if line.is_empty() || !line.starts_with("data: ") {
            return None;
        }

        let data = &line[6..]; // Remove "data: " prefix
        
        if data == "[DONE]" {
            return Some(Ok(ProviderEvent::Complete { token_usage: None }));
        }

        match serde_json::from_str::<StreamChunk>(data) {
            Ok(chunk) => {
                if let Some(choice) = chunk.choices.first() {
                    if let Some(content) = &choice.delta.content {
                        return Some(Ok(ProviderEvent::ContentChunk {
                            content: content.clone(),
                        }));
                    }
                    
                    if let Some(tool_calls) = &choice.delta.tool_calls {
                        if let Some(first_call) = tool_calls.first() {
                            return Some(Ok(ProviderEvent::ToolCall {
                                name: first_call.function.name.clone(),
                                parameters: serde_json::from_str(&first_call.function.arguments)
                                    .unwrap_or(serde_json::Value::Null),
                            }));
                        }
                    }
                }
                None
            }
            Err(e) => Some(Err(ProviderError::ResponseParsing(format!(
                "Failed to parse stream chunk: {}",
                e
            )))),
        }
    }
}

#[async_trait]
impl Provider for OpenAICompatibleProvider {
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError> {
        let request = ChatCompletionRequest {
            model: self.model.id.clone(),
            messages: self.convert_messages(messages),
            max_tokens: self.max_tokens,
            temperature: self.temperature,
            stream: false,
            tools: if tools.is_empty() || !self.supports_tools {
                None
            } else {
                Some(self.convert_tools(tools))
            },
            tool_choice: None,
        };

        let mut request_builder = self
            .client
            .post(&format!("{}/chat/completions", self.base_url))
            .header("Content-Type", "application/json");

        // Add authorization header if API key is provided
        if !self.api_key.is_empty() {
            request_builder = request_builder.header("Authorization", format!("Bearer {}", self.api_key));
        }

        let response = request_builder
            .json(&request)
            .send()
            .await
            .map_err(|e| ProviderError::NetworkError(e.to_string()))?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            return Err(ProviderError::ApiError(format!(
                "{} API error {}: {}",
                self.provider_name, status, error_text
            )));
        }

        let completion: ChatCompletionResponse = response
            .json()
            .await
            .map_err(|e| ProviderError::ResponseParsing(e.to_string()))?;

        let choice = completion
            .choices
            .first()
            .ok_or_else(|| ProviderError::ResponseParsing("No choices in response".into()))?;

        let token_usage = completion.usage.clone().map(|u| TokenUsage {
            input_tokens: u.prompt_tokens,
            output_tokens: u.completion_tokens,
            total_tokens: u.total_tokens,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        });

        Ok(ProviderResponse {
            content: choice.message.content.clone(),
            tool_calls: choice
                .message
                .tool_calls
                .as_ref()
                .map(|calls| self.convert_tool_calls(calls))
                .unwrap_or_default(),
            token_usage,
            metadata: serde_json::to_value(&completion).unwrap_or_default(),
        })
    }

    async fn stream_response(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError> {
        if !self.supports_streaming {
            return Err(ProviderError::Unavailable(format!(
                "{} does not support streaming",
                self.provider_name
            )));
        }

        let (tx, rx) = mpsc::channel(100);

        let request = ChatCompletionRequest {
            model: self.model.id.clone(),
            messages: self.convert_messages(messages),
            max_tokens: self.max_tokens,
            temperature: self.temperature,
            stream: true,
            tools: if tools.is_empty() || !self.supports_tools {
                None
            } else {
                Some(self.convert_tools(tools))
            },
            tool_choice: None,
        };

        let client = self.client.clone();
        let base_url = self.base_url.clone();
        let api_key = self.api_key.clone();
        let provider_name = self.provider_name.clone();

        tokio::spawn(async move {
            let mut request_builder = client
                .post(&format!("{}/chat/completions", base_url))
                .header("Content-Type", "application/json");

            // Add authorization header if API key is provided
            if !api_key.is_empty() {
                request_builder = request_builder.header("Authorization", format!("Bearer {}", api_key));
            }

            let response = match request_builder
                .json(&request)
                .send()
                .await
            {
                Ok(response) => response,
                Err(e) => {
                    let _ = tx.send(Err(ProviderError::NetworkError(e.to_string()))).await;
                    return;
                }
            };

            if !response.status().is_success() {
                let status = response.status();
                let error_text = response.text().await.unwrap_or_default();
                let _ = tx
                    .send(Err(ProviderError::ApiError(format!(
                        "{} API error {}: {}",
                        provider_name, status, error_text
                    ))))
                    .await;
                return;
            }

            let mut stream = response.bytes_stream();
            let mut buffer = String::new();

            while let Some(chunk) = stream.next().await {
                match chunk {
                    Ok(bytes) => {
                        buffer.push_str(&String::from_utf8_lossy(&bytes));
                        
                        // Process complete lines
                        while let Some(line_end) = buffer.find('\n') {
                            let line = buffer[..line_end].trim().to_string();
                            buffer = buffer[line_end + 1..].to_string();
                            
                            if let Some(event) = OpenAICompatibleProvider::parse_stream_line(&OpenAICompatibleProvider {
                                client: client.clone(),
                                api_key: api_key.clone(),
                                base_url: base_url.clone(),
                                model: Model {
                                    id: "temp".to_string(),
                                    name: "temp".to_string(),
                                    provider: provider_name.clone(),
                                    context_length: 0,
                                    max_output_tokens: 0,
                                    supports_tools: false,
                                    supports_streaming: false,
                                    supports_vision: false,
                                    cost_per_input_token: 0.0,
                                    cost_per_output_token: 0.0,
                                    capabilities: ModelCapabilities::default(),
                                },
                                max_tokens: None,
                                temperature: None,
                                provider_name: provider_name.clone(),
                                supports_tools: false,
                                supports_streaming: false,
                            }, &line) {
                                if let Err(_) = tx.send(event).await {
                                    break;
                                }
                            }
                        }
                    }
                    Err(e) => {
                        let _ = tx.send(Err(ProviderError::NetworkError(e.to_string()))).await;
                        break;
                    }
                }
            }
        });

        Ok(rx)
    }

    fn model(&self) -> &Model {
        &self.model
    }

    fn name(&self) -> &str {
        &self.provider_name
    }

    async fn is_available(&self) -> bool {
        // Simple health check - try to make a minimal request
        let test_request = ChatCompletionRequest {
            model: self.model.id.clone(),
            messages: vec![OpenAIMessage {
                role: "user".to_string(),
                content: "test".to_string(),
                tool_calls: None,
                tool_call_id: None,
            }],
            max_tokens: Some(1),
            temperature: Some(0.0),
            stream: false,
            tools: None,
            tool_choice: None,
        };

        let mut request_builder = self
            .client
            .post(&format!("{}/chat/completions", self.base_url))
            .header("Content-Type", "application/json");

        if !self.api_key.is_empty() {
            request_builder = request_builder.header("Authorization", format!("Bearer {}", self.api_key));
        }

        request_builder
            .json(&test_request)
            .send()
            .await
            .map(|response| response.status().is_success())
            .unwrap_or(false)
    }
}