ralph-coder 0.2.1

use crate::errors::{RalphError, Result};
use crate::providers::{
    ContentPart, LlmProvider, LlmResponse, Message, MessageContent, Role, StopReason, ToolCall,
    ToolDef,
};
use async_trait::async_trait;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;

pub const DEFAULT_MODEL: &str = "gpt-4o";
const CONTEXT_WINDOW: u64 = 128_000;
const API_URL: &str = "https://api.openai.com/v1/chat/completions";
const MAX_TOKENS: u32 = 4096;

pub struct OpenAiProvider {
    api_key: String,
    model: String,
    base_url: String,
    client: reqwest::Client,
}

impl OpenAiProvider {
    pub fn new(api_key: String, model: Option<String>, base_url: Option<String>) -> Self {
        Self {
            api_key,
            model: model.unwrap_or_else(|| DEFAULT_MODEL.to_string()),
            base_url: base_url.unwrap_or_else(|| API_URL.to_string()),
            client: reqwest::Client::new(),
        }
    }
}

// ── Request / response types ──────────────────────────────────────────────────

#[derive(Serialize)]
struct OpenAiRequest {
    model: String,
    messages: Vec<OaiMessage>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<OaiTool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_completion_tokens: Option<u32>,
}

#[derive(Serialize)]
struct OaiMessage {
    role: String,
    content: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OaiToolCallOut>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
}

#[derive(Serialize, Deserialize)]
struct OaiToolCallOut {
    id: String,
    #[serde(rename = "type")]
    kind: String,
    function: OaiFunctionCall,
}

#[derive(Serialize, Deserialize)]
struct OaiFunctionCall {
    name: String,
    arguments: String,
}

#[derive(Serialize)]
struct OaiTool {
    #[serde(rename = "type")]
    kind: String,
    function: OaiToolFunction,
}

#[derive(Serialize)]
struct OaiToolFunction {
    name: String,
    description: String,
    parameters: Value,
}

#[derive(Deserialize)]
struct OaiResponse {
    choices: Vec<OaiChoice>,
    usage: OaiUsage,
}

#[derive(Deserialize)]
struct OaiChoice {
    message: OaiRespMessage,
    finish_reason: Option<String>,
}

#[derive(Deserialize)]
struct OaiRespMessage {
    content: Option<String>,
    tool_calls: Option<Vec<OaiToolCallOut>>,
}

#[derive(Deserialize)]
struct OaiUsage {
    #[serde(default)]
    prompt_tokens: u64,
    #[serde(default)]
    completion_tokens: u64,
    total_tokens: u64,
}

// ── SSE streaming types ───────────────────────────────────────────────────────

#[derive(Deserialize)]
struct OaiStreamChunk {
    choices: Vec<OaiStreamChoice>,
    #[serde(default)]
    usage: Option<OaiUsage>,
}

#[derive(Deserialize)]
struct OaiStreamChoice {
    delta: OaiStreamDelta,
    finish_reason: Option<String>,
}

#[derive(Deserialize, Default)]
struct OaiStreamDelta {
    content: Option<String>,
    tool_calls: Option<Vec<OaiToolCallChunk>>,
}

#[derive(Deserialize)]
struct OaiToolCallChunk {
    index: usize,
    id: Option<String>,
    function: Option<OaiFunctionChunk>,
}

#[derive(Deserialize)]
struct OaiFunctionChunk {
    name: Option<String>,
    arguments: Option<String>,
}

fn messages_to_oai(messages: &[Message]) -> Vec<OaiMessage> {
    messages
        .iter()
        .map(|m| {
            let role = match m.role {
                Role::System => "system",
                Role::User => "user",
                Role::Assistant => "assistant",
                Role::Tool => "tool",
            };

            // Assistant messages that contain tool-use parts must have the
            // `tool_calls` array set — otherwise the API rejects subsequent
            // tool-result messages ("must be a response to a preceding message
            // with tool_calls").
            if matches!(m.role, Role::Assistant) {
                if let MessageContent::Parts(parts) = &m.content {
                    let text: String = parts
                        .iter()
                        .filter_map(|p| match p {
                            ContentPart::Text { text } => Some(text.as_str()),
                            _ => None,
                        })
                        .collect::<Vec<_>>()
                        .join("\n");

                    let tool_calls: Vec<OaiToolCallOut> = parts
                        .iter()
                        .filter_map(|p| match p {
                            ContentPart::ToolUse { id, name, input } => Some(OaiToolCallOut {
                                id: id.clone(),
                                kind: "function".to_string(),
                                function: OaiFunctionCall {
                                    name: name.clone(),
                                    arguments: serde_json::to_string(input)
                                        .unwrap_or_else(|_| "{}".to_string()),
                                },
                            }),
                            _ => None,
                        })
                        .collect();

                    return OaiMessage {
                        role: "assistant".to_string(),
                        content: if text.is_empty() {
                            json!(null)
                        } else {
                            json!(text)
                        },
                        tool_calls: if tool_calls.is_empty() {
                            None
                        } else {
                            Some(tool_calls)
                        },
                        tool_call_id: None,
                        name: None,
                    };
                }
            }

            OaiMessage {
                role: role.to_string(),
                content: json!(m.content.as_text()),
                tool_calls: None,
                tool_call_id: m.tool_call_id.clone(),
                name: m.name.clone(),
            }
        })
        .collect()
}

#[async_trait]
impl LlmProvider for OpenAiProvider {
    async fn chat(&self, messages: &[Message], tools: &[ToolDef]) -> Result<LlmResponse> {
        let oai_tools: Vec<OaiTool> = tools
            .iter()
            .map(|t| OaiTool {
                kind: "function".to_string(),
                function: OaiToolFunction {
                    name: t.name.clone(),
                    description: t.description.clone(),
                    parameters: t.parameters.clone(),
                },
            })
            .collect();

        // search-preview models use max_completion_tokens; others use max_tokens
        let is_search_model = self.model.contains("search");
        let body = OpenAiRequest {
            model: self.model.clone(),
            messages: messages_to_oai(messages),
            tools: oai_tools,
            max_tokens: if is_search_model {
                None
            } else {
                Some(MAX_TOKENS)
            },
            max_completion_tokens: if is_search_model {
                Some(MAX_TOKENS)
            } else {
                None
            },
        };

        let resp = self
            .client
            .post(&self.base_url)
            .bearer_auth(&self.api_key)
            .json(&body)
            .send()
            .await?;

        let status = resp.status();
        if status == 401 {
            return Err(RalphError::LlmAuth {
                provider: "openai".to_string(),
            });
        }
        if status == 429 {
            return Err(RalphError::LlmRateLimit {
                provider: "openai".to_string(),
                attempts: 1,
            });
        }
        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(RalphError::LlmApi {
                provider: "openai".to_string(),
                message: format!("HTTP {}: {}", status, body),
            });
        }

        let parsed: OaiResponse = resp
            .json()
            .await
            .map_err(|e| RalphError::LlmResponseParse(e.to_string()))?;

        let choice =
            parsed.choices.into_iter().next().ok_or_else(|| {
                RalphError::LlmResponseParse("No choices in response".to_string())
            })?;

        let stop_reason = match choice.finish_reason.as_deref() {
            Some("tool_calls") => StopReason::ToolUse,
            Some("stop") => StopReason::Stop,
            Some("length") => StopReason::MaxTokens,
            _ => StopReason::EndTurn,
        };

        let tool_calls = choice
            .message
            .tool_calls
            .unwrap_or_default()
            .into_iter()
            .map(|tc| {
                let args: Value = serde_json::from_str(&tc.function.arguments)
                    .unwrap_or(Value::Object(Default::default()));
                ToolCall {
                    id: tc.id,
                    name: tc.function.name,
                    arguments: args,
                }
            })
            .collect();

        Ok(LlmResponse {
            text: choice.message.content,
            tool_calls,
            input_tokens: parsed.usage.prompt_tokens,
            output_tokens: parsed.usage.completion_tokens,
            reasoning_tokens: 0,
            reasoning_content: None,
            tokens_used: parsed.usage.total_tokens,
            stop_reason,
        })
    }

    async fn chat_streaming(
        &self,
        messages: &[Message],
        tools: &[ToolDef],
        token_tx: &tokio::sync::mpsc::UnboundedSender<String>,
    ) -> Result<LlmResponse> {
        let oai_tools: Vec<OaiTool> = tools
            .iter()
            .map(|t| OaiTool {
                kind: "function".to_string(),
                function: OaiToolFunction {
                    name: t.name.clone(),
                    description: t.description.clone(),
                    parameters: t.parameters.clone(),
                },
            })
            .collect();

        let is_search_model = self.model.contains("search");
        let mut body = json!({
            "model": self.model,
            "messages": messages_to_oai(messages),
            "stream": true,
            "stream_options": { "include_usage": true },
        });
        if !oai_tools.is_empty() {
            body["tools"] = json!(oai_tools);
        }
        if !is_search_model {
            body["max_tokens"] = json!(MAX_TOKENS);
        } else {
            body["max_completion_tokens"] = json!(MAX_TOKENS);
        }

        let resp = self
            .client
            .post(&self.base_url)
            .bearer_auth(&self.api_key)
            .json(&body)
            .send()
            .await?;

        let status = resp.status();
        if status == 401 {
            return Err(RalphError::LlmAuth {
                provider: "openai".to_string(),
            });
        }
        if status == 429 {
            return Err(RalphError::LlmRateLimit {
                provider: "openai".to_string(),
                attempts: 1,
            });
        }
        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(RalphError::LlmApi {
                provider: "openai".to_string(),
                message: format!("HTTP {}: {}", status, body),
            });
        }

        // Process SSE stream
        let mut stream = resp.bytes_stream();
        let mut buf = String::new();
        let mut text_parts: Vec<String> = Vec::new();
        // index → (id, name, accumulated_args)
        let mut tool_chunks: HashMap<usize, (String, String, String)> = HashMap::new();
        let mut total_tokens: u64 = 0;
        let mut finish_reason: Option<String> = None;

        while let Some(chunk) = stream.next().await {
            let bytes = chunk.map_err(|e| RalphError::LlmApi {
                provider: "openai".to_string(),
                message: e.to_string(),
            })?;
            buf.push_str(&String::from_utf8_lossy(&bytes));

            // Process every complete line in the buffer
            loop {
                let Some(pos) = buf.find('\n') else { break };
                let line = buf[..pos].trim().to_string();
                buf.drain(..pos + 1);

                if !line.starts_with("data: ") {
                    continue;
                }
                let data = line[6..].trim();
                if data == "[DONE]" {
                    break;
                }

                let Ok(parsed) = serde_json::from_str::<OaiStreamChunk>(data) else {
                    continue;
                };
                if let Some(usage) = parsed.usage {
                    total_tokens = usage.total_tokens;
                }
                for choice in parsed.choices {
                    if let Some(r) = choice.finish_reason {
                        finish_reason = Some(r);
                    }
                    if let Some(content) = choice.delta.content {
                        if !content.is_empty() {
                            let _ = token_tx.send(content.clone());
                            text_parts.push(content);
                        }
                    }
                    if let Some(tcs) = choice.delta.tool_calls {
                        for tc in tcs {
                            let entry = tool_chunks
                                .entry(tc.index)
                                .or_insert_with(|| (String::new(), String::new(), String::new()));
                            if let Some(id) = tc.id {
                                entry.0 = id;
                            }
                            if let Some(func) = tc.function {
                                if let Some(name) = func.name {
                                    entry.1 = name;
                                }
                                if let Some(args) = func.arguments {
                                    entry.2.push_str(&args);
                                }
                            }
                        }
                    }
                }
            }
        }

        let text = if text_parts.is_empty() {
            None
        } else {
            Some(text_parts.join(""))
        };

        let mut sorted: Vec<_> = tool_chunks.into_iter().collect();
        sorted.sort_by_key(|(i, _)| *i);
        let tool_calls = sorted
            .into_iter()
            .map(|(_, (id, name, args))| {
                let arguments: Value =
                    serde_json::from_str(&args).unwrap_or(Value::Object(Default::default()));
                ToolCall {
                    id,
                    name,
                    arguments,
                }
            })
            .collect();

        let stop_reason = match finish_reason.as_deref() {
            Some("tool_calls") => StopReason::ToolUse,
            Some("stop") => StopReason::Stop,
            Some("length") => StopReason::MaxTokens,
            _ => StopReason::EndTurn,
        };

        Ok(LlmResponse {
            text,
            tool_calls,
            input_tokens: 0,
            output_tokens: 0,
            reasoning_tokens: 0,
            reasoning_content: None,
            tokens_used: total_tokens,
            stop_reason,
        })
    }

    fn supports_streaming(&self) -> bool {
        true
    }

    fn name(&self) -> &str {
        "openai"
    }
    fn context_window(&self) -> u64 {
        CONTEXT_WINDOW
    }
    fn default_model(&self) -> &str {
        DEFAULT_MODEL
    }
}