coderlib 0.1.0

//! AWS Bedrock provider implementation for CoderLib
//!
//! This module provides integration with AWS Bedrock's foundation models including
//! Claude, Llama, Titan, and other models available through Bedrock.
//! Supports both streaming and non-streaming responses with comprehensive error handling.

use async_trait::async_trait;
use aws_config::{BehaviorVersion, Region};
use aws_sdk_bedrockruntime::{Client as BedrockClient, types::ContentBlock};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio::sync::mpsc;

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

/// AWS Bedrock provider for foundation models
pub struct BedrockProvider {
    client: BedrockClient,
    model: Model,
    region: String,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
}

/// Bedrock request format for Claude models
#[derive(Debug, Serialize)]
struct BedrockClaudeRequest {
    anthropic_version: String,
    max_tokens: u32,
    messages: Vec<BedrockMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<String>,
}

/// Bedrock message format
#[derive(Debug, Serialize)]
struct BedrockMessage {
    role: String,
    content: String,
}

/// Bedrock response format for Claude models
#[derive(Debug, Deserialize)]
struct BedrockClaudeResponse {
    content: Vec<BedrockContent>,
    usage: BedrockUsage,
    stop_reason: Option<String>,
}

/// Bedrock content block
#[derive(Debug, Deserialize)]
struct BedrockContent {
    #[serde(rename = "type")]
    content_type: String,
    text: Option<String>,
}

/// Bedrock usage information
#[derive(Debug, Deserialize)]
struct BedrockUsage {
    input_tokens: u32,
    output_tokens: u32,
}

/// Bedrock streaming response chunk
#[derive(Debug, Deserialize)]
struct BedrockStreamChunk {
    #[serde(rename = "type")]
    chunk_type: String,
    delta: Option<BedrockDelta>,
    usage: Option<BedrockUsage>,
}

/// Bedrock streaming delta
#[derive(Debug, Deserialize)]
struct BedrockDelta {
    #[serde(rename = "type")]
    delta_type: String,
    text: Option<String>,
}

impl BedrockProvider {
    /// Create a new Bedrock provider
    pub async fn new(
        region: Option<String>,
        model: Model,
        max_tokens: Option<u32>,
        temperature: Option<f32>,
    ) -> Result<Self, ProviderError> {
        // Load AWS configuration
        let config = aws_config::defaults(BehaviorVersion::latest())
            .region(Region::new(region.clone().unwrap_or_else(|| "us-east-1".to_string())))
            .load()
            .await;
        
        let client = BedrockClient::new(&config);
        
        Ok(Self {
            client,
            model,
            region: region.unwrap_or_else(|| "us-east-1".to_string()),
            max_tokens,
            temperature,
        })
    }
    
    /// Convert CoderLib messages to Bedrock format
    fn convert_messages(&self, messages: Vec<Message>) -> (Vec<BedrockMessage>, Option<String>) {
        let mut bedrock_messages = Vec::new();
        let mut system_message = None;
        
        for message in messages {
            match message.role.as_str() {
                "system" => {
                    system_message = Some(message.content);
                }
                "user" | "assistant" => {
                    bedrock_messages.push(BedrockMessage {
                        role: message.role,
                        content: message.content,
                    });
                }
                _ => {
                    // Skip unknown roles
                }
            }
        }
        
        (bedrock_messages, system_message)
    }
    
    /// Create request payload for Claude models
    fn create_claude_request(&self, messages: Vec<Message>) -> Result<String, ProviderError> {
        let (bedrock_messages, system) = self.convert_messages(messages);
        
        let request = BedrockClaudeRequest {
            anthropic_version: "bedrock-2023-05-31".to_string(),
            max_tokens: self.max_tokens.unwrap_or(4000),
            messages: bedrock_messages,
            temperature: self.temperature,
            system,
        };
        
        serde_json::to_string(&request)
            .map_err(|e| ProviderError::RequestFailed(format!("Failed to serialize request: {}", e)))
    }
    
    /// Parse Claude response
    fn parse_claude_response(&self, response_body: &str) -> Result<ProviderResponse, ProviderError> {
        let response: BedrockClaudeResponse = serde_json::from_str(response_body)
            .map_err(|e| ProviderError::ResponseParseFailed(format!("Failed to parse response: {}", e)))?;
        
        let content = response.content
            .into_iter()
            .filter_map(|c| c.text)
            .collect::<Vec<_>>()
            .join("");
        
        Ok(ProviderResponse {
            content,
            tool_calls: Vec::new(), // TODO: Implement tool calls for Bedrock
            token_usage: Some(TokenUsage {
                input_tokens: response.usage.input_tokens as u64,
                output_tokens: response.usage.output_tokens as u64,
                cache_creation_tokens: 0,
                cache_read_tokens: 0,
            }),
            metadata: json!({
                "stop_reason": response.stop_reason,
                "model": self.model.id,
                "provider": "bedrock"
            }),
        })
    }
    
    /// Determine model family for request formatting
    fn get_model_family(&self) -> &str {
        if self.model.id.contains("claude") {
            "claude"
        } else if self.model.id.contains("llama") {
            "llama"
        } else if self.model.id.contains("titan") {
            "titan"
        } else {
            "claude" // Default to Claude format
        }
    }
}

#[async_trait]
impl Provider for BedrockProvider {
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        _tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError> {
        let model_family = self.get_model_family();
        
        match model_family {
            "claude" => {
                let request_body = self.create_claude_request(messages)?;
                
                let response = self.client
                    .invoke_model()
                    .model_id(&self.model.id)
                    .content_type("application/json")
                    .body(aws_sdk_bedrockruntime::primitives::Blob::new(request_body.as_bytes()))
                    .send()
                    .await
                    .map_err(|e| ProviderError::RequestFailed(format!("Bedrock request failed: {}", e)))?;
                
                let response_body = String::from_utf8(response.body().as_ref().to_vec())
                    .map_err(|e| ProviderError::ResponseParseFailed(format!("Invalid UTF-8 response: {}", e)))?;
                
                self.parse_claude_response(&response_body)
            }
            _ => {
                Err(ProviderError::ModelNotSupported(format!("Model family '{}' not yet supported", model_family)))
            }
        }
    }
    
    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 model_family = self.get_model_family();
        
        match model_family {
            "claude" => {
                let request_body = self.create_claude_request(messages)?;
                let client = self.client.clone();
                let model_id = self.model.id.clone();
                
                tokio::spawn(async move {
                    let result = client
                        .invoke_model_with_response_stream()
                        .model_id(&model_id)
                        .content_type("application/json")
                        .body(aws_sdk_bedrockruntime::primitives::Blob::new(request_body.as_bytes()))
                        .send()
                        .await;
                    
                    match result {
                        Ok(mut response) => {
                            let mut stream = response.body;
                            
                            while let Some(event) = stream.recv().await {
                                match event {
                                    Ok(event) => {
                                        if let Some(chunk) = event.chunk {
                                            let chunk_data = String::from_utf8_lossy(chunk.bytes());
                                            
                                            if let Ok(parsed_chunk) = serde_json::from_str::<BedrockStreamChunk>(&chunk_data) {
                                                if let Some(delta) = parsed_chunk.delta {
                                                    if let Some(text) = delta.text {
                                                        let _ = tx.send(Ok(ProviderEvent::Delta(text))).await;
                                                    }
                                                }
                                                
                                                if parsed_chunk.chunk_type == "message_stop" {
                                                    let _ = tx.send(Ok(ProviderEvent::Done)).await;
                                                    break;
                                                }
                                            }
                                        }
                                    }
                                    Err(e) => {
                                        let _ = tx.send(Err(ProviderError::StreamingFailed(format!("Stream error: {}", e)))).await;
                                        break;
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(ProviderError::RequestFailed(format!("Bedrock streaming request failed: {}", e)))).await;
                        }
                    }
                });
            }
            _ => {
                let _ = tx.send(Err(ProviderError::ModelNotSupported(format!("Streaming not supported for model family '{}'", model_family)))).await;
            }
        }
        
        Ok(rx)
    }
    
    fn model(&self) -> &Model {
        &self.model
    }
    
    fn name(&self) -> &str {
        "bedrock"
    }
    
    async fn is_available(&self) -> bool {
        // Test availability by making a simple request to list foundation models
        self.client
            .list_foundation_models()
            .send()
            .await
            .is_ok()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::ModelCapabilities;
    
    #[test]
    fn test_message_conversion() {
        let model = Model {
            id: "anthropic.claude-3-sonnet-20240229-v1:0".to_string(),
            name: "Claude 3 Sonnet".to_string(),
            provider: "bedrock".to_string(),
            context_length: 200000,
            max_output_tokens: 4000,
            supports_tools: false,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.003,
            cost_per_output_token: 0.015,
            capabilities: ModelCapabilities::default(),
        };
        
        let provider = BedrockProvider {
            client: BedrockClient::new(&aws_config::SdkConfig::builder().build()),
            model,
            region: "us-east-1".to_string(),
            max_tokens: Some(1000),
            temperature: Some(0.7),
        };
        
        let messages = vec![
            Message {
                role: "system".to_string(),
                content: "You are a helpful assistant.".to_string(),
            },
            Message {
                role: "user".to_string(),
                content: "Hello!".to_string(),
            },
        ];
        
        let (bedrock_messages, system) = provider.convert_messages(messages);
        
        assert_eq!(bedrock_messages.len(), 1);
        assert_eq!(bedrock_messages[0].role, "user");
        assert_eq!(bedrock_messages[0].content, "Hello!");
        assert_eq!(system, Some("You are a helpful assistant.".to_string()));
    }
    
    #[test]
    fn test_model_family_detection() {
        let claude_model = Model {
            id: "anthropic.claude-3-sonnet-20240229-v1:0".to_string(),
            name: "Claude 3 Sonnet".to_string(),
            provider: "bedrock".to_string(),
            context_length: 200000,
            max_output_tokens: 4000,
            supports_tools: false,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.003,
            cost_per_output_token: 0.015,
            capabilities: ModelCapabilities::default(),
        };
        
        let provider = BedrockProvider {
            client: BedrockClient::new(&aws_config::SdkConfig::builder().build()),
            model: claude_model,
            region: "us-east-1".to_string(),
            max_tokens: Some(1000),
            temperature: Some(0.7),
        };
        
        assert_eq!(provider.get_model_family(), "claude");
    }
}