//! OpenAI chat completion API -- blocking and streaming.

use anyhow::{Context, Result};
use futures_core::Stream;
use reqwest::header;
use serde::{Deserialize, Serialize};

const API_URL: &str = "https://api.openai.com/v1/chat/completions";

/// OpenAI chat completion client.
pub struct ChatGptClient {
    client: reqwest::Client,
    model: String,
}

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

// -- Request types -----------------------------------------------------------

#[derive(Debug, Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "std::ops::Not::not")]
    stream: bool,
}

// -- Non-streaming response types --------------------------------------------

#[derive(Debug, Deserialize)]
struct ChatResponse {
    choices: Vec<Choice>,
}

#[derive(Debug, Deserialize)]
struct Choice {
    message: Message,
}

// -- Streaming response types ------------------------------------------------

#[derive(Debug, Deserialize)]
struct StreamChunk {
    choices: Vec<StreamChoice>,
}

#[derive(Debug, Deserialize)]
struct StreamChoice {
    delta: Delta,
}

#[derive(Debug, Deserialize)]
struct Delta {
    content: Option<String>,
}

// -- Errors ------------------------------------------------------------------

#[derive(Debug, thiserror::Error)]
pub enum LlmError {
    #[error("OpenAI API error (HTTP {status}): {body}")]
    Api { status: u16, body: String },

    #[error(transparent)]
    Transport(#[from] reqwest::Error),

    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

impl ChatGptClient {
    /// Create a new client. The API key is baked into the underlying
    /// `reqwest::Client` as a default header so it doesn't need to be
    /// cloned per-request.
    pub fn new(api_key: &str, model: &str) -> Result<Self> {
        let mut headers = header::HeaderMap::new();
        let mut auth = header::HeaderValue::from_str(&format!("Bearer {api_key}"))
            .context("invalid API key characters")?;
        auth.set_sensitive(true);
        headers.insert(header::AUTHORIZATION, auth);

        let client = reqwest::Client::builder()
            .default_headers(headers)
            .build()
            .context("failed to build HTTP client")?;

        Ok(Self {
            client,
            model: model.to_owned(),
        })
    }

    /// Send a blocking chat completion request, returning the full response.
    pub async fn chat(&self, messages: Vec<Message>) -> Result<String, LlmError> {
        let request = ChatRequest {
            model: self.model.clone(),
            messages,
            temperature: None,
            stream: false,
        };

        let response = self.client.post(API_URL).json(&request).send().await?;
        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            return Err(LlmError::Api {
                status: status.as_u16(),
                body,
            });
        }

        let parsed: ChatResponse = response.json().await?;
        Ok(parsed
            .choices
            .into_iter()
            .next()
            .map(|c| c.message.content)
            .unwrap_or_default())
    }

    /// Stream a chat completion, yielding content tokens as they arrive.
    ///
    /// The returned stream is `'static` -- it clones the HTTP client and model
    /// name so callers don't need to worry about lifetimes.
    pub fn chat_stream(
        &self,
        messages: Vec<Message>,
    ) -> impl Stream<Item = Result<String, LlmError>> + Send {
        let client = self.client.clone();
        let model = self.model.clone();

        async_stream::try_stream! {
            let request = ChatRequest {
                model,
                messages,
                temperature: None,
                stream: true,
            };

            let mut response = client.post(API_URL).json(&request).send().await?;
            if !response.status().is_success() {
                let status = response.status().as_u16();
                // Read error body via chunk() to avoid .text() consuming by value.
                let mut body = String::new();
                while let Some(chunk) = response.chunk().await? {
                    body.push_str(&String::from_utf8_lossy(&chunk));
                }
                Err(LlmError::Api { status, body })?;
            }
            let mut buffer = String::new();

            while let Some(chunk) = response.chunk().await? {
                buffer.push_str(&String::from_utf8_lossy(&chunk));

                // Process complete SSE events (delimited by double newline).
                while let Some(pos) = buffer.find("\n\n") {
                    let event = buffer[..pos].to_owned();
                    buffer = buffer[pos + 2..].to_owned();

                    for line in event.lines() {
                        let data = match line.strip_prefix("data: ") {
                            Some(d) => d.trim(),
                            None => continue,
                        };

                        if data == "[DONE]" {
                            return;
                        }

                        if let Ok(parsed) = serde_json::from_str::<StreamChunk>(data) {
                            for choice in parsed.choices {
                                if let Some(content) = choice.delta.content {
                                    yield content;
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}
