//! Anthropic Claude API client implementation

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::env;

use super::provider::{
    LlmProvider, LlmRequest, LlmResponse, LlmStream, MessageRole, StreamChunk, StreamMetadata,
    TokenUsage,
};
use crate::error::{Error, Result};
use futures::stream::StreamExt;

const CLAUDE_API_URL: &str = "https://api.anthropic.com/v1/messages";
const CLAUDE_API_VERSION: &str = "2023-06-01";

/// Claude API client
pub struct ClaudeProvider {
    api_key: String,
    client: reqwest::Client,
}

/// Claude API request format
#[derive(Debug, Serialize)]
struct ClaudeRequest {
    model: String,
    messages: Vec<ClaudeMessage>,
    max_tokens: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
}

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

/// Claude API response format
#[derive(Debug, Deserialize)]
struct ClaudeResponse {
    #[allow(dead_code)]
    id: String,
    #[allow(dead_code)]
    #[serde(rename = "type")]
    response_type: String,
    #[allow(dead_code)]
    role: String,
    content: Vec<ClaudeContent>,
    model: String,
    #[allow(dead_code)]
    stop_reason: Option<String>,
    usage: ClaudeUsage,
}

#[derive(Debug, Deserialize)]
struct ClaudeContent {
    #[serde(rename = "type")]
    content_type: String,
    text: String,
}

#[derive(Debug, Deserialize)]
struct ClaudeUsage {
    input_tokens: usize,
    output_tokens: usize,
}

impl ClaudeProvider {
    /// Create a new Claude provider
    pub fn new() -> Result<Self> {
        let api_key = env::var("ANTHROPIC_API_KEY").map_err(|_| {
            Error::execution(
                "ANTHROPIC_API_KEY environment variable not set",
                "Set ANTHROPIC_API_KEY to use Claude models",
            )
        })?;

        let client = reqwest::Client::new();

        Ok(Self { api_key, client })
    }

    /// Create a provider with explicit API key
    pub fn with_api_key(api_key: String) -> Self {
        let client = reqwest::Client::new();
        Self { api_key, client }
    }
}

#[async_trait]
impl LlmProvider for ClaudeProvider {
    async fn complete(&self, request: LlmRequest) -> Result<LlmResponse> {
        // Convert generic messages to Claude format
        let messages: Vec<ClaudeMessage> = request
            .messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System) // System handled separately
            .map(|m| ClaudeMessage {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "assistant".to_string(),
                    MessageRole::System => "user".to_string(), // Fallback
                },
                content: m.content,
            })
            .collect();

        // Build Claude request
        let claude_request = ClaudeRequest {
            model: request.model.clone(),
            messages,
            max_tokens: request.max_tokens.unwrap_or(4096),
            temperature: request.temperature,
            system: request.system,
            stream: None,
        };

        // Make API call
        let response = self
            .client
            .post(CLAUDE_API_URL)
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", CLAUDE_API_VERSION)
            .header("content-type", "application/json")
            .json(&claude_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call Claude API", e.to_string()))?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Claude API error: {}", status),
                error_text,
            ));
        }

        // Parse response
        let claude_response: ClaudeResponse = response
            .json()
            .await
            .map_err(|e| Error::execution("Failed to parse Claude API response", e.to_string()))?;

        // Extract text content
        let content = claude_response
            .content
            .into_iter()
            .filter_map(|c| {
                if c.content_type == "text" {
                    Some(c.text)
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
            .join("\n");

        Ok(LlmResponse {
            content,
            model: claude_response.model,
            usage: Some(TokenUsage {
                input_tokens: claude_response.usage.input_tokens,
                output_tokens: claude_response.usage.output_tokens,
                total_tokens: claude_response.usage.input_tokens
                    + claude_response.usage.output_tokens,
            }),
        })
    }

    async fn stream(&self, request: LlmRequest) -> Result<LlmStream> {
        use futures::stream;

        // Convert generic messages to Claude format
        let messages: Vec<ClaudeMessage> = request
            .messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System)
            .map(|m| ClaudeMessage {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "assistant".to_string(),
                    MessageRole::System => "user".to_string(),
                },
                content: m.content,
            })
            .collect();

        // Build Claude streaming request
        let claude_request = ClaudeRequest {
            model: request.model.clone(),
            messages,
            max_tokens: request.max_tokens.unwrap_or(4096),
            temperature: request.temperature,
            system: request.system,
            stream: Some(true),
        };

        // Make streaming API call
        let response = self
            .client
            .post(CLAUDE_API_URL)
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", CLAUDE_API_VERSION)
            .header("content-type", "application/json")
            .json(&claude_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call Claude API", e.to_string()))?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Claude API error: {}", status),
                error_text,
            ));
        }

        // Convert response to SSE stream
        use futures::TryStreamExt;
        let byte_stream = response.bytes_stream();
        let model = request.model.clone();

        let sse_stream = byte_stream
            .map_err(|e| Error::execution("Stream error", e.to_string()))
            .map(move |chunk_result| {
                match chunk_result {
                    Ok(bytes) => {
                        // Parse SSE events
                        let text = String::from_utf8_lossy(&bytes);
                        parse_sse_events(&text, &model)
                    }
                    Err(e) => vec![Err(e)],
                }
            })
            .flat_map(stream::iter);

        Ok(Box::pin(sse_stream))
    }

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

/// Parse SSE events from Claude streaming response
fn parse_sse_events(text: &str, model: &str) -> Vec<Result<StreamChunk>> {
    let mut chunks = Vec::new();

    for line in text.lines() {
        if let Some(data) = line.strip_prefix("data: ") {
            // Parse JSON event
            if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
                match event.get("type").and_then(|t| t.as_str()) {
                    Some("content_block_delta") => {
                        // Extract text delta
                        if let Some(delta) = event.get("delta") {
                            if let Some(text) = delta.get("text").and_then(|t| t.as_str()) {
                                chunks.push(Ok(StreamChunk::Content(text.to_string())));
                            }
                        }
                    }
                    Some("message_stop") => {
                        // Stream completed - extract usage if available
                        let usage = event.get("usage").and_then(|u| {
                            Some(TokenUsage {
                                input_tokens: u.get("input_tokens")?.as_u64()? as usize,
                                output_tokens: u.get("output_tokens")?.as_u64()? as usize,
                                total_tokens: (u.get("input_tokens")?.as_u64()?
                                    + u.get("output_tokens")?.as_u64()?)
                                    as usize,
                            })
                        });

                        chunks.push(Ok(StreamChunk::Done(StreamMetadata {
                            model: model.to_string(),
                            usage,
                        })));
                    }
                    Some("message_delta") => {
                        // Extract usage from delta
                        if let Some(usage_obj) = event.get("usage") {
                            let usage = Some(TokenUsage {
                                input_tokens: usage_obj
                                    .get("input_tokens")
                                    .and_then(|v| v.as_u64())
                                    .unwrap_or(0)
                                    as usize,
                                output_tokens: usage_obj
                                    .get("output_tokens")
                                    .and_then(|v| v.as_u64())
                                    .unwrap_or(0)
                                    as usize,
                                total_tokens: (usage_obj
                                    .get("input_tokens")
                                    .and_then(|v| v.as_u64())
                                    .unwrap_or(0)
                                    + usage_obj
                                        .get("output_tokens")
                                        .and_then(|v| v.as_u64())
                                        .unwrap_or(0))
                                    as usize,
                            });

                            chunks.push(Ok(StreamChunk::Done(StreamMetadata {
                                model: model.to_string(),
                                usage,
                            })));
                        }
                    }
                    Some("error") => {
                        let error_msg = event
                            .get("error")
                            .and_then(|e| e.get("message"))
                            .and_then(|m| m.as_str())
                            .unwrap_or("Unknown error");
                        chunks.push(Ok(StreamChunk::Error(error_msg.to_string())));
                    }
                    _ => {}
                }
            }
        }
    }

    chunks
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::LlmMessage;

    #[test]
    fn test_create_provider_without_api_key() {
        // Should fail if ANTHROPIC_API_KEY not set
        let original_key = env::var("ANTHROPIC_API_KEY").ok();
        env::remove_var("ANTHROPIC_API_KEY");

        let result = ClaudeProvider::new();
        assert!(result.is_err());

        // Restore original key if it existed
        if let Some(key) = original_key {
            env::set_var("ANTHROPIC_API_KEY", key);
        }
    }

    #[test]
    fn test_create_provider_with_explicit_key() {
        let provider = ClaudeProvider::with_api_key("test-key".to_string());
        assert_eq!(provider.api_key, "test-key");
        assert_eq!(provider.name(), "Claude");
    }

    #[test]
    fn test_parse_sse_content_block_delta() {
        let sse_data =
            r#"data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"Hello"}}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected Content chunk"),
        }
    }

    #[test]
    fn test_parse_sse_message_stop() {
        let sse_data = r#"data: {"type":"message_stop"}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Done(metadata)) => {
                assert_eq!(metadata.model, model);
                assert!(metadata.usage.is_none());
            }
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_sse_message_delta_with_usage() {
        let sse_data =
            r#"data: {"type":"message_delta","usage":{"input_tokens":100,"output_tokens":50}}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Done(metadata)) => {
                assert_eq!(metadata.model, model);
                let usage = metadata.usage.as_ref().unwrap();
                assert_eq!(usage.input_tokens, 100);
                assert_eq!(usage.output_tokens, 50);
                assert_eq!(usage.total_tokens, 150);
            }
            _ => panic!("Expected Done chunk with usage"),
        }
    }

    #[test]
    fn test_parse_sse_error() {
        let sse_data = r#"data: {"type":"error","error":{"message":"Rate limit exceeded"}}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Error(msg)) => assert_eq!(msg, "Rate limit exceeded"),
            _ => panic!("Expected Error chunk"),
        }
    }

    #[test]
    fn test_parse_sse_multiple_events() {
        let sse_data = r#"data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"Hello"}}
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":" world"}}
data: {"type":"message_delta","usage":{"input_tokens":10,"output_tokens":5}}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 3);

        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected first Content chunk"),
        }

        match &chunks[1] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, " world"),
            _ => panic!("Expected second Content chunk"),
        }

        match &chunks[2] {
            Ok(StreamChunk::Done(metadata)) => {
                assert!(metadata.usage.is_some());
            }
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_sse_ignores_unknown_events() {
        let sse_data = r#"data: {"type":"unknown_event"}
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"test"}}"#;
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        // Should only parse the content delta, ignore unknown
        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "test"),
            _ => panic!("Expected Content chunk"),
        }
    }

    #[test]
    fn test_parse_sse_empty_input() {
        let sse_data = "";
        let model = "claude-3-5-haiku-20241022";

        let chunks = parse_sse_events(sse_data, model);

        assert_eq!(chunks.len(), 0);
    }

    #[tokio::test]
    #[ignore] // Only run with real API key
    async fn test_claude_api_call() {
        let provider = ClaudeProvider::new().expect("ANTHROPIC_API_KEY must be set");

        let request = LlmRequest {
            model: "claude-3-5-haiku-20241022".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Say hello in exactly 3 words.".to_string(),
            }],
            max_tokens: Some(50),
            temperature: Some(0.7),
            system: None,
        };

        let response = provider.complete(request).await;
        assert!(response.is_ok());

        let response = response.unwrap();
        assert!(!response.content.is_empty());
        assert!(response.usage.is_some());

        println!("Response: {}", response.content);
        println!("Usage: {:?}", response.usage);
    }
}
