coderlib 0.1.0

//! Google Cloud Vertex AI provider implementation for CoderLib
//!
//! This module provides integration with Google Cloud Vertex AI, supporting
//! Gemini models and other Google AI models through the Vertex AI platform.

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

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

/// Google Cloud Vertex AI provider
pub struct VertexProvider {
    client: Client,
    project_id: String,
    location: String,
    access_token: String,
    model: Model,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
}

/// Vertex AI request format
#[derive(Debug, Serialize)]
struct VertexRequest {
    contents: Vec<VertexContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    generation_config: Option<VertexGenerationConfig>,
    #[serde(skip_serializing_if = "Option::is_none")]
    safety_settings: Option<Vec<VertexSafetySetting>>,
}

/// Vertex AI content
#[derive(Debug, Serialize)]
struct VertexContent {
    role: String,
    parts: Vec<VertexPart>,
}

/// Vertex AI content part
#[derive(Debug, Serialize)]
struct VertexPart {
    text: String,
}

/// Vertex AI generation configuration
#[derive(Debug, Serialize)]
struct VertexGenerationConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_output_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_k: Option<u32>,
}

/// Vertex AI safety setting
#[derive(Debug, Serialize)]
struct VertexSafetySetting {
    category: String,
    threshold: String,
}

/// Vertex AI response format
#[derive(Debug, Deserialize)]
struct VertexResponse {
    candidates: Vec<VertexCandidate>,
    #[serde(rename = "usageMetadata")]
    usage_metadata: Option<VertexUsageMetadata>,
}

/// Vertex AI candidate
#[derive(Debug, Deserialize)]
struct VertexCandidate {
    content: VertexResponseContent,
    #[serde(rename = "finishReason")]
    finish_reason: Option<String>,
}

/// Vertex AI response content
#[derive(Debug, Deserialize)]
struct VertexResponseContent {
    parts: Vec<VertexResponsePart>,
    role: String,
}

/// Vertex AI response part
#[derive(Debug, Deserialize)]
struct VertexResponsePart {
    text: Option<String>,
}

/// Vertex AI usage metadata
#[derive(Debug, Deserialize)]
struct VertexUsageMetadata {
    #[serde(rename = "promptTokenCount")]
    prompt_token_count: Option<u32>,
    #[serde(rename = "candidatesTokenCount")]
    candidates_token_count: Option<u32>,
    #[serde(rename = "totalTokenCount")]
    total_token_count: Option<u32>,
}

impl VertexProvider {
    /// Create a new Vertex AI provider
    pub fn new(
        project_id: String,
        location: String,
        access_token: String,
        model: Model,
        max_tokens: Option<u32>,
        temperature: Option<f32>,
    ) -> Result<Self, ProviderError> {
        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(120))
            .build()
            .map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;
        
        Ok(Self {
            client,
            project_id,
            location,
            access_token,
            model,
            max_tokens,
            temperature,
        })
    }
    
    /// Get the API endpoint URL
    fn get_api_url(&self) -> String {
        format!(
            "https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:generateContent",
            self.location, self.project_id, self.location, self.model.id
        )
    }
    
    /// Get the streaming API endpoint URL
    fn get_streaming_api_url(&self) -> String {
        format!(
            "https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:streamGenerateContent",
            self.location, self.project_id, self.location, self.model.id
        )
    }
    
    /// Convert CoderLib messages to Vertex format
    fn convert_messages(&self, messages: Vec<Message>) -> Vec<VertexContent> {
        messages
            .into_iter()
            .filter_map(|msg| {
                // Vertex AI uses "user" and "model" roles
                let role = match msg.role {
                    MessageRole::User => "user",
                    MessageRole::Assistant => "model",
                    MessageRole::System => "user", // System messages are treated as user messages
                    MessageRole::Tool => return None, // Skip tool messages for now
                };

                Some(VertexContent {
                    role: role.to_string(),
                    parts: vec![VertexPart {
                        text: msg.content.as_text(),
                    }],
                })
            })
            .collect()
    }
    
    /// Create request payload
    fn create_request(&self, messages: Vec<Message>) -> VertexRequest {
        let generation_config = if self.max_tokens.is_some() || self.temperature.is_some() {
            Some(VertexGenerationConfig {
                temperature: self.temperature,
                max_output_tokens: self.max_tokens,
                top_p: None,
                top_k: None,
            })
        } else {
            None
        };
        
        // Default safety settings (permissive for development)
        let safety_settings = Some(vec![
            VertexSafetySetting {
                category: "HARM_CATEGORY_HARASSMENT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            VertexSafetySetting {
                category: "HARM_CATEGORY_HATE_SPEECH".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            VertexSafetySetting {
                category: "HARM_CATEGORY_SEXUALLY_EXPLICIT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            VertexSafetySetting {
                category: "HARM_CATEGORY_DANGEROUS_CONTENT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
        ]);
        
        VertexRequest {
            contents: self.convert_messages(messages),
            generation_config,
            safety_settings,
        }
    }
    
    /// Parse response
    fn parse_response(&self, response: VertexResponse) -> Result<ProviderResponse, ProviderError> {
        let candidate = response.candidates
            .into_iter()
            .next()
            .ok_or_else(|| ProviderError::ResponseParsing("No candidates in response".to_string()))?;
        
        let content = candidate.content.parts
            .into_iter()
            .filter_map(|part| part.text)
            .collect::<Vec<_>>()
            .join("");
        
        let token_usage = response.usage_metadata.map(|usage| TokenUsage {
            input_tokens: usage.prompt_token_count.unwrap_or(0),
            output_tokens: usage.candidates_token_count.unwrap_or(0),
            total_tokens: usage.prompt_token_count.unwrap_or(0) + usage.candidates_token_count.unwrap_or(0),
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        });
        
        Ok(ProviderResponse {
            content,
            tool_calls: Vec::new(), // TODO: Implement tool calls for Vertex AI
            token_usage,
            metadata: serde_json::json!({
                "finish_reason": candidate.finish_reason,
                "model": self.model.id,
                "provider": "vertex",
                "project_id": self.project_id,
                "location": self.location
            }),
        })
    }
}

#[async_trait]
impl Provider for VertexProvider {
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        _tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError> {
        let request = self.create_request(messages);
        
        let response = self.client
            .post(&self.get_api_url())
            .header("Authorization", format!("Bearer {}", self.access_token))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .map_err(|e| ProviderError::Http(e))?;
        
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            return Err(ProviderError::ApiError(format!(
                "Vertex AI error {}: {}", status, error_text
            )));
        }
        
        let vertex_response: VertexResponse = response.json().await
            .map_err(|e| ProviderError::ResponseParsing(format!("Failed to parse Vertex AI response: {}", e)))?;
        
        self.parse_response(vertex_response)
    }
    
    async fn stream_response(
        &self,
        messages: Vec<Message>,
        _tools: Vec<Box<dyn Tool>>,
    ) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError> {
        let (tx, rx) = mpsc::channel(100);
        
        let request = self.create_request(messages);
        let client = self.client.clone();
        let url = self.get_streaming_api_url();
        let access_token = self.access_token.clone();
        
        tokio::spawn(async move {
            let response_result = client
                .post(&url)
                .header("Authorization", format!("Bearer {}", access_token))
                .header("Content-Type", "application/json")
                .json(&request)
                .send()
                .await;
            
            match response_result {
                Ok(response) => {
                    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!(
                            "Vertex AI error {}: {}", status, error_text
                        )))).await;
                        return;
                    }
                    
                    // Vertex AI streaming uses newline-delimited JSON
                    let text = response.text().await.unwrap_or_default();
                    
                    for line in text.lines() {
                        if line.trim().is_empty() {
                            continue;
                        }
                        
                        if let Ok(chunk) = serde_json::from_str::<VertexResponse>(line) {
                            if let Some(candidate) = chunk.candidates.first() {
                                for part in &candidate.content.parts {
                                    if let Some(text) = &part.text {
                                        let _ = tx.send(Ok(ProviderEvent::ContentChunk { content: text.clone() })).await;
                                    }
                                }
                                
                                if candidate.finish_reason.is_some() {
                                    let _ = tx.send(Ok(ProviderEvent::Complete { token_usage: None })).await;
                                    return;
                                }
                            }
                        }
                    }
                    
                    let _ = tx.send(Ok(ProviderEvent::Complete { token_usage: None })).await;
                }
                Err(e) => {
                    let _ = tx.send(Err(ProviderError::Http(e))).await;
                }
            }
        });
        
        Ok(rx)
    }
    
    fn model(&self) -> &Model {
        &self.model
    }
    
    fn name(&self) -> &str {
        "vertex"
    }
    
    async fn is_available(&self) -> bool {
        // Test availability with a simple request
        let test_messages = vec![Message {
            id: "test".to_string(),
            session_id: "test".to_string(),
            role: MessageRole::User,
            content: crate::storage::MessageContent::Text("test".to_string()),
            timestamp: chrono::Utc::now(),
            metadata: serde_json::Value::Null,
        }];
        
        let request = self.create_request(test_messages);
        
        self.client
            .post(&self.get_api_url())
            .header("Authorization", format!("Bearer {}", self.access_token))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .map(|response| response.status().is_success())
            .unwrap_or(false)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::ModelCapabilities;
    
    #[test]
    fn test_vertex_provider_creation() {
        let model = Model {
            id: "gemini-1.5-pro".to_string(),
            name: "Gemini 1.5 Pro".to_string(),
            provider: "vertex".to_string(),
            context_length: 1000000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.0035,
            cost_per_output_token: 0.0105,
            capabilities: ModelCapabilities::default(),
        };
        
        let provider = VertexProvider::new(
            "test-project".to_string(),
            "us-central1".to_string(),
            "test-token".to_string(),
            model,
            Some(1000),
            Some(0.7),
        ).unwrap();
        
        assert_eq!(provider.name(), "vertex");
        assert_eq!(provider.project_id, "test-project");
        assert_eq!(provider.location, "us-central1");
    }
    
    #[test]
    fn test_api_url_generation() {
        let model = Model {
            id: "gemini-1.5-pro".to_string(),
            name: "Gemini 1.5 Pro".to_string(),
            provider: "vertex".to_string(),
            context_length: 1000000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.0035,
            cost_per_output_token: 0.0105,
            capabilities: ModelCapabilities::default(),
        };
        
        let provider = VertexProvider::new(
            "my-project".to_string(),
            "europe-west1".to_string(),
            "test-token".to_string(),
            model,
            None,
            None,
        ).unwrap();
        
        let expected_url = "https://europe-west1-aiplatform.googleapis.com/v1/projects/my-project/locations/europe-west1/publishers/google/models/gemini-1.5-pro:generateContent";
        assert_eq!(provider.get_api_url(), expected_url);
    }
    
    #[test]
    fn test_message_conversion() {
        let model = Model {
            id: "gemini-1.5-pro".to_string(),
            name: "Gemini 1.5 Pro".to_string(),
            provider: "vertex".to_string(),
            context_length: 1000000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.0035,
            cost_per_output_token: 0.0105,
            capabilities: ModelCapabilities::default(),
        };
        
        let provider = VertexProvider::new(
            "test-project".to_string(),
            "us-central1".to_string(),
            "test-token".to_string(),
            model,
            None,
            None,
        ).unwrap();
        
        let messages = vec![
            Message {
                id: "1".to_string(),
                session_id: "test".to_string(),
                role: MessageRole::System,
                content: crate::storage::MessageContent::Text("You are a helpful assistant.".to_string()),
                timestamp: chrono::Utc::now(),
                metadata: serde_json::Value::Null,
            },
            Message {
                id: "2".to_string(),
                session_id: "test".to_string(),
                role: MessageRole::User,
                content: crate::storage::MessageContent::Text("Hello!".to_string()),
                timestamp: chrono::Utc::now(),
                metadata: serde_json::Value::Null,
            },
            Message {
                id: "3".to_string(),
                session_id: "test".to_string(),
                role: MessageRole::Assistant,
                content: crate::storage::MessageContent::Text("Hi there!".to_string()),
                timestamp: chrono::Utc::now(),
                metadata: serde_json::Value::Null,
            },
        ];
        
        let vertex_contents = provider.convert_messages(messages);
        
        assert_eq!(vertex_contents.len(), 3);
        assert_eq!(vertex_contents[0].role, "user"); // system -> user
        assert_eq!(vertex_contents[1].role, "user");
        assert_eq!(vertex_contents[2].role, "model"); // assistant -> model
    }
}