coderlib 0.1.0

//! Anthropic provider implementation for CoderLib
//!
//! This module provides integration with Anthropic's Claude models including
//! Claude 3.5 Sonnet, Claude 3 Haiku, and Claude 3 Opus with support for
//! streaming, tool use, and large context windows.

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;

/// Anthropic provider for Claude models
#[derive(Debug, Clone)]
pub struct AnthropicProvider {
    client: Client,
    api_key: String,
    base_url: String,
    model: Model,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
}

/// Anthropic message request
#[derive(Debug, Serialize)]
struct MessageRequest {
    model: String,
    max_tokens: u32,
    messages: Vec<AnthropicMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<AnthropicTool>>,
}

/// Anthropic message format
#[derive(Debug, Serialize, Deserialize)]
struct AnthropicMessage {
    role: String,
    content: AnthropicContent,
}

/// Anthropic content can be text or array of content blocks
#[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
enum AnthropicContent {
    Text(String),
    Blocks(Vec<ContentBlock>),
}

/// Anthropic content block
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
enum ContentBlock {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "tool_use")]
    ToolUse {
        id: String,
        name: String,
        input: serde_json::Value,
    },
    #[serde(rename = "tool_result")]
    ToolResult {
        tool_use_id: String,
        content: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        is_error: Option<bool>,
    },
}

/// Anthropic tool definition
#[derive(Debug, Serialize)]
struct AnthropicTool {
    name: String,
    description: String,
    input_schema: serde_json::Value,
}

/// Anthropic message response
#[derive(Debug, Serialize, Deserialize)]
struct MessageResponse {
    id: String,
    #[serde(rename = "type")]
    response_type: String,
    role: String,
    content: Vec<ContentBlock>,
    model: String,
    stop_reason: Option<String>,
    stop_sequence: Option<String>,
    usage: Usage,
}

/// Anthropic usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Usage {
    input_tokens: u32,
    output_tokens: u32,
}

/// Anthropic streaming event
#[derive(Debug, Deserialize)]
struct StreamEvent {
    #[serde(rename = "type")]
    event_type: String,
    #[serde(flatten)]
    data: serde_json::Value,
}

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

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

        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(120)) // Anthropic can be slower
            .build()
            .map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;

        Ok(Self {
            client,
            api_key,
            base_url: base_url.unwrap_or_else(|| "https://api.anthropic.com".to_string()),
            model,
            max_tokens,
            temperature,
        })
    }

    /// Convert CoderLib messages to Anthropic format
    fn convert_messages(&self, messages: Vec<Message>) -> (Option<String>, Vec<AnthropicMessage>) {
        let mut system_message = None;
        let mut anthropic_messages = Vec::new();

        for msg in messages {
            match msg.role {
                crate::storage::MessageRole::System => {
                    // Anthropic handles system messages separately
                    system_message = Some(msg.content.as_text());
                }
                crate::storage::MessageRole::User => {
                    anthropic_messages.push(AnthropicMessage {
                        role: "user".to_string(),
                        content: self.convert_content(&msg.content),
                    });
                }
                crate::storage::MessageRole::Assistant => {
                    anthropic_messages.push(AnthropicMessage {
                        role: "assistant".to_string(),
                        content: self.convert_content(&msg.content),
                    });
                }
                crate::storage::MessageRole::Tool => {
                    // Tool results are handled as part of user messages in Anthropic
                    if let crate::storage::MessageContent::ToolResult { result, tool_call_id } = &msg.content {
                        anthropic_messages.push(AnthropicMessage {
                            role: "user".to_string(),
                            content: AnthropicContent::Blocks(vec![ContentBlock::ToolResult {
                                tool_use_id: tool_call_id.clone(),
                                content: result.clone(),
                                is_error: None,
                            }]),
                        });
                    }
                }
            }
        }

        (system_message, anthropic_messages)
    }

    /// Convert message content to Anthropic format
    fn convert_content(&self, content: &crate::storage::MessageContent) -> AnthropicContent {
        match content {
            crate::storage::MessageContent::Text(text) => AnthropicContent::Text(text.clone()),
            crate::storage::MessageContent::Structured(parts) => {
                let blocks: Vec<ContentBlock> = parts
                    .iter()
                    .filter_map(|part| match part {
                        crate::storage::ContentPart::Text { text } => {
                            Some(ContentBlock::Text { text: text.clone() })
                        }
                        crate::storage::ContentPart::Code { code, language } => {
                            Some(ContentBlock::Text {
                                text: format!("```{}\n{}\n```", language, code),
                            })
                        }
                        crate::storage::ContentPart::ToolCall { name, parameters } => {
                            Some(ContentBlock::ToolUse {
                                id: uuid::Uuid::new_v4().to_string(),
                                name: name.clone(),
                                input: parameters.clone(),
                            })
                        }
                        _ => None, // Skip unsupported content types
                    })
                    .collect();
                
                if blocks.len() == 1 {
                    if let ContentBlock::Text { text } = &blocks[0] {
                        return AnthropicContent::Text(text.clone());
                    }
                }
                
                AnthropicContent::Blocks(blocks)
            }
            crate::storage::MessageContent::ToolCall { name, parameters, id } => {
                AnthropicContent::Blocks(vec![ContentBlock::ToolUse {
                    id: id.clone(),
                    name: name.clone(),
                    input: parameters.clone(),
                }])
            }
            crate::storage::MessageContent::ToolResult { result, tool_call_id } => {
                AnthropicContent::Blocks(vec![ContentBlock::ToolResult {
                    tool_use_id: tool_call_id.clone(),
                    content: result.clone(),
                    is_error: None,
                }])
            }
        }
    }

    /// Convert tools to Anthropic format
    fn convert_tools(&self, tools: Vec<Box<dyn Tool>>) -> Vec<AnthropicTool> {
        tools
            .into_iter()
            .map(|tool| AnthropicTool {
                name: tool.name().to_string(),
                description: tool.description().to_string(),
                input_schema: tool.parameter_schema(),
            })
            .collect()
    }

    /// Convert Anthropic tool uses to CoderLib format
    fn convert_tool_calls(&self, content: &[ContentBlock]) -> Vec<ProviderToolCall> {
        content
            .iter()
            .filter_map(|block| match block {
                ContentBlock::ToolUse { id, name, input } => Some(ProviderToolCall {
                    id: id.clone(),
                    name: name.clone(),
                    parameters: input.clone(),
                }),
                _ => None,
            })
            .collect()
    }

    /// Parse streaming event
    fn parse_stream_event(&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::<StreamEvent>(data) {
            Ok(event) => {
                match event.event_type.as_str() {
                    "content_block_delta" => {
                        if let Ok(delta) = serde_json::from_value::<StreamDelta>(event.data) {
                            if let Some(text) = delta.text {
                                return Some(Ok(ProviderEvent::ContentChunk { content: text }));
                            }
                        }
                    }
                    "message_stop" => {
                        return Some(Ok(ProviderEvent::Complete { token_usage: None }));
                    }
                    "error" => {
                        let error_msg = event.data.get("message")
                            .and_then(|v| v.as_str())
                            .unwrap_or("Unknown error");
                        return Some(Ok(ProviderEvent::Error {
                            error: error_msg.to_string(),
                        }));
                    }
                    _ => {} // Ignore other event types
                }
                None
            }
            Err(e) => Some(Err(ProviderError::ResponseParsing(format!(
                "Failed to parse stream event: {}",
                e
            )))),
        }
    }
}

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

        let response = self
            .client
            .post(&format!("{}/v1/messages", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .header("anthropic-version", "2023-06-01")
            .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!(
                "Anthropic API error {}: {}",
                status, error_text
            )));
        }

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

        // Extract text content and tool calls
        let mut content_text = String::new();
        let tool_calls = self.convert_tool_calls(&message_response.content);

        for block in &message_response.content {
            if let ContentBlock::Text { text } = block {
                if !content_text.is_empty() {
                    content_text.push('\n');
                }
                content_text.push_str(text);
            }
        }

        Ok(ProviderResponse {
            content: content_text,
            tool_calls,
            token_usage: Some(TokenUsage {
                input_tokens: message_response.usage.input_tokens,
                output_tokens: message_response.usage.output_tokens,
                total_tokens: message_response.usage.input_tokens + message_response.usage.output_tokens,
                cache_creation_tokens: 0,
                cache_read_tokens: 0,
            }),
            metadata: serde_json::to_value(&message_response).unwrap_or_default(),
        })
    }

    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 (system, anthropic_messages) = self.convert_messages(messages);

        let request = MessageRequest {
            model: self.model.id.clone(),
            max_tokens: self.max_tokens.unwrap_or(4000),
            messages: anthropic_messages,
            system,
            temperature: self.temperature,
            stream: true,
            tools: if tools.is_empty() {
                None
            } else {
                Some(self.convert_tools(tools))
            },
        };

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

        tokio::spawn(async move {
            let response = match client
                .post(&format!("{}/v1/messages", base_url))
                .header("Authorization", format!("Bearer {}", api_key))
                .header("Content-Type", "application/json")
                .header("anthropic-version", "2023-06-01")
                .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!(
                        "Anthropic API error {}: {}",
                        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) = AnthropicProvider::parse_stream_event(&AnthropicProvider {
                                client: client.clone(),
                                api_key: api_key.clone(),
                                base_url: base_url.clone(),
                                model: Model {
                                    id: "temp".to_string(),
                                    name: "temp".to_string(),
                                    provider: "anthropic".to_string(),
                                    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,
                            }, &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 {
        "anthropic"
    }

    async fn is_available(&self) -> bool {
        // Simple health check - try to make a minimal request
        let request = MessageRequest {
            model: "claude-3-haiku-20240307".to_string(),
            max_tokens: 1,
            messages: vec![AnthropicMessage {
                role: "user".to_string(),
                content: AnthropicContent::Text("test".to_string()),
            }],
            system: None,
            temperature: Some(0.0),
            stream: false,
            tools: None,
        };

        self.client
            .post(&format!("{}/v1/messages", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .header("anthropic-version", "2023-06-01")
            .json(&request)
            .send()
            .await
            .map(|response| response.status().is_success())
            .unwrap_or(false)
    }
}