//! OpenAI 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 OPENAI_API_URL: &str = "https://api.openai.com/v1/chat/completions";

/// OpenAI API client
pub struct OpenAIProvider {
    api_key: String,
    client: reqwest::Client,
}

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

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

/// OpenAI API response format
#[derive(Debug, Deserialize)]
struct OpenAIResponse {
    #[allow(dead_code)]
    id: String,
    #[allow(dead_code)]
    object: String,
    #[allow(dead_code)]
    created: u64,
    model: String,
    choices: Vec<OpenAIChoice>,
    usage: OpenAIUsage,
}

#[derive(Debug, Deserialize)]
struct OpenAIChoice {
    #[allow(dead_code)]
    index: usize,
    message: OpenAIMessage,
    #[allow(dead_code)]
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct OpenAIUsage {
    prompt_tokens: usize,
    completion_tokens: usize,
    total_tokens: usize,
}

impl OpenAIProvider {
    /// Create a new OpenAI provider
    pub fn new() -> Result<Self> {
        let api_key = env::var("OPENAI_API_KEY").map_err(|_| {
            Error::execution(
                "OPENAI_API_KEY environment variable not set",
                "Set OPENAI_API_KEY to use OpenAI 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 OpenAIProvider {
    async fn complete(&self, request: LlmRequest) -> Result<LlmResponse> {
        // Convert generic messages to OpenAI format
        let mut messages: Vec<OpenAIMessage> = request
            .messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System)
            .map(|m| OpenAIMessage {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "assistant".to_string(),
                    MessageRole::System => "system".to_string(),
                },
                content: m.content,
            })
            .collect();

        // Add system message at the beginning if provided
        if let Some(system_prompt) = request.system {
            messages.insert(
                0,
                OpenAIMessage {
                    role: "system".to_string(),
                    content: system_prompt,
                },
            );
        }

        // Build OpenAI request
        let openai_request = OpenAIRequest {
            model: request.model.clone(),
            messages,
            max_tokens: request.max_tokens,
            temperature: request.temperature,
            stream: None,
        };

        // Make API call
        let response = self
            .client
            .post(OPENAI_API_URL)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&openai_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call OpenAI 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!("OpenAI API error: {}", status),
                error_text,
            ));
        }

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

        // Extract content from first choice
        let content = openai_response
            .choices
            .first()
            .map(|choice| choice.message.content.clone())
            .unwrap_or_default();

        Ok(LlmResponse {
            content,
            model: openai_response.model,
            usage: Some(TokenUsage {
                input_tokens: openai_response.usage.prompt_tokens,
                output_tokens: openai_response.usage.completion_tokens,
                total_tokens: openai_response.usage.total_tokens,
            }),
        })
    }

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

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

        // Add system message at the beginning if provided
        if let Some(system_prompt) = request.system {
            messages.insert(
                0,
                OpenAIMessage {
                    role: "system".to_string(),
                    content: system_prompt,
                },
            );
        }

        // Build OpenAI streaming request
        let openai_request = OpenAIRequest {
            model: request.model.clone(),
            messages,
            max_tokens: request.max_tokens,
            temperature: request.temperature,
            stream: Some(true),
        };

        // Make streaming API call
        let response = self
            .client
            .post(OPENAI_API_URL)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&openai_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call OpenAI 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!("OpenAI 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 {
        "OpenAI"
    }
}

/// Parse SSE events from OpenAI 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: ") {
            // Check for completion signal
            if data.trim() == "[DONE]" {
                chunks.push(Ok(StreamChunk::Done(StreamMetadata {
                    model: model.to_string(),
                    usage: None, // OpenAI doesn't send usage in streaming mode
                })));
                continue;
            }

            // Parse JSON event
            if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
                // Extract delta content
                if let Some(choices) = event.get("choices").and_then(|c| c.as_array()) {
                    if let Some(choice) = choices.first() {
                        if let Some(delta) = choice.get("delta") {
                            if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                                chunks.push(Ok(StreamChunk::Content(content.to_string())));
                            }
                        }
                    }
                }

                // Check for errors
                if let Some(error) = event.get("error") {
                    let error_msg = error
                        .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 OPENAI_API_KEY not set
        let original_key = env::var("OPENAI_API_KEY").ok();
        env::remove_var("OPENAI_API_KEY");

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

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

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

    #[test]
    fn test_parse_sse_content_delta() {
        let sse_data = r#"data: {"choices":[{"delta":{"content":"Hello"}}]}"#;
        let model = "gpt-4";

        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_done_signal() {
        let sse_data = "data: [DONE]";
        let model = "gpt-4";

        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_error() {
        let sse_data = r#"data: {"error":{"message":"Rate limit exceeded"}}"#;
        let model = "gpt-4";

        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: {"choices":[{"delta":{"content":"Hello"}}]}
data: {"choices":[{"delta":{"content":" world"}}]}
data: [DONE]"#;
        let model = "gpt-4";

        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(_)) => {}
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_sse_empty_delta() {
        let sse_data = r#"data: {"choices":[{"delta":{}}]}"#;
        let model = "gpt-4";

        let chunks = parse_sse_events(sse_data, model);

        // Should not produce any chunks for empty delta
        assert_eq!(chunks.len(), 0);
    }

    #[test]
    fn test_parse_sse_empty_input() {
        let sse_data = "";
        let model = "gpt-4";

        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_openai_api_call() {
        let provider = OpenAIProvider::new().expect("OPENAI_API_KEY must be set");

        let request = LlmRequest {
            model: "gpt-4o-mini".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);
    }
}
