//! OpenAI provider impl.
//!
//! Wraps the Chat Completions API and supports the GPT-4o / o1 / o3
//! families plus Azure OpenAI via a custom base URL. JSON mode is a
//! first-class flag (`response_format`); tool use mirrors the OpenAI
//! tools array. Token counting is local via `tiktoken-rs`: `cl100k_base`
//! for the GPT-4 family and `o200k_base` for GPT-4o and the o-series.

use std::time::Duration;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::{json, Value};
use tiktoken_rs::CoreBPE;

use crate::eval::error::ProviderError;
use crate::eval::provider::{
    map_reqwest_error, Capabilities, CompletionOptions, Cost, FinishReason, LlmProvider, Message,
    ProviderId, Response, ResponseContent, RetryPolicy, Role, ToolCall, ToolChoice, Usage,
};
use crate::eval::providers::parse_embedding;
use crate::eval::providers::retry::run_with_retry;

const DEFAULT_BASE_URL: &str = "https://api.openai.com";

/// OpenAI provider client. Covers Chat Completions across GPT-4o,
/// GPT-4o-mini, o1-preview, o1-mini, and o3-mini.
pub struct OpenAiProvider {
    client: reqwest::Client,
    api_key: String,
    organization: Option<String>,
    base_url: String,
    retry: RetryPolicy,
    request_timeout: Duration,
}

/// Config bag for [`OpenAiProvider::new`].
#[derive(Debug, Clone)]
pub struct OpenAiConfig {
    /// Resolved API key.
    pub api_key: String,
    /// Optional organization header (`OPENAI_ORG_ID`).
    pub organization: Option<String>,
    /// Override the base URL (Azure OpenAI, wiremock, etc.).
    pub base_url: Option<String>,
    /// HTTP request timeout. Default 30s.
    pub request_timeout: Duration,
    /// HTTP / HTTPS proxy policy.
    pub proxy: crate::network::ProxyConfig,
}

impl Default for OpenAiConfig {
    fn default() -> Self {
        Self {
            api_key: String::new(),
            organization: None,
            base_url: None,
            request_timeout: Duration::from_secs(30),
            proxy: crate::network::ProxyConfig::default(),
        }
    }
}

impl OpenAiProvider {
    /// Build a new client. Returns [`ProviderError::Transport`] when the
    /// `reqwest` builder fails to initialize TLS.
    pub fn new(config: OpenAiConfig) -> Result<Self, ProviderError> {
        let builder = reqwest::Client::builder().timeout(config.request_timeout);
        let builder = config
            .proxy
            .apply(builder)
            .map_err(|e| ProviderError::Transport(e.to_string()))?;
        let client = builder
            .build()
            .map_err(|e| ProviderError::Transport(e.to_string()))?;
        Ok(Self {
            client,
            api_key: config.api_key,
            organization: config.organization,
            base_url: config.base_url.unwrap_or_else(|| DEFAULT_BASE_URL.into()),
            retry: RetryPolicy::default(),
            request_timeout: config.request_timeout,
        })
    }
}

#[async_trait]
impl LlmProvider for OpenAiProvider {
    fn name(&self) -> ProviderId {
        ProviderId::Openai
    }

    fn supported_capabilities(&self) -> Capabilities {
        Capabilities {
            json_mode: true,
            tool_use: true,
            streaming: true,
            structured_output: true,
            // Local tiktoken counts are exact for OpenAI's BPE families.
            token_count_exact: true,
        }
    }

    async fn chat_completion(
        &self,
        messages: &[Message],
        options: &CompletionOptions,
    ) -> Result<Response, ProviderError> {
        let body = build_chat_body(messages, options);
        let url = format!("{}/v1/chat/completions", self.base_url);
        let response_json: Value = run_with_retry(&self.retry, || async {
            let mut req = self
                .client
                .post(&url)
                .bearer_auth(&self.api_key)
                .json(&body);
            if let Some(org) = &self.organization {
                req = req.header("OpenAI-Organization", org);
            }
            let resp = req
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_chat_response(response_json, &options.model)
    }

    async fn count_tokens(&self, text: &str, model: &str) -> Result<u32, ProviderError> {
        let bpe = load_encoding_for(model)?;
        let count = bpe.encode_with_special_tokens(text).len();
        Ok(count as u32)
    }

    async fn embed(&self, text: &str, model: &str) -> Result<Vec<f32>, ProviderError> {
        let url = format!("{}/v1/embeddings", self.base_url);
        let body = json!({ "model": model, "input": text });
        let response_json: Value = run_with_retry(&self.retry, || async {
            let mut req = self
                .client
                .post(&url)
                .bearer_auth(&self.api_key)
                .json(&body);
            if let Some(org) = &self.organization {
                req = req.header("OpenAI-Organization", org);
            }
            let resp = req
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_embedding(&response_json, "/data/0/embedding")
    }

    fn estimate_cost(&self, usage: &Usage, model: &str) -> Cost {
        let rates = openai_rates(model);
        Cost {
            input_dollars: rates.input_per_token * f64::from(usage.input_tokens),
            output_dollars: rates.output_per_token * f64::from(usage.output_tokens),
            cache_read_dollars: None,
            cache_write_dollars: None,
        }
    }
}

fn load_encoding_for(model: &str) -> Result<CoreBPE, ProviderError> {
    let lower = model.to_ascii_lowercase();
    let result =
        if lower.starts_with("gpt-4o") || lower.starts_with("o1") || lower.starts_with("o3") {
            tiktoken_rs::o200k_base()
        } else {
            tiktoken_rs::cl100k_base()
        };
    result.map_err(|e| ProviderError::Transport(format!("tiktoken: {e}")))
}

#[derive(Debug, Clone, Copy)]
struct OpenAiRates {
    input_per_token: f64,
    output_per_token: f64,
}

/// Published per-token rates for the common OpenAI models. Unknown
/// models fall back to GPT-4o-mini, the cheapest mainline option.
fn openai_rates(model: &str) -> OpenAiRates {
    let lower = model.to_ascii_lowercase();
    if lower.starts_with("gpt-4o-mini") {
        OpenAiRates {
            input_per_token: 0.15 / 1_000_000.0,
            output_per_token: 0.60 / 1_000_000.0,
        }
    } else if lower.starts_with("gpt-4o") {
        OpenAiRates {
            input_per_token: 2.50 / 1_000_000.0,
            output_per_token: 10.00 / 1_000_000.0,
        }
    } else if lower.starts_with("o3-mini") {
        OpenAiRates {
            input_per_token: 1.10 / 1_000_000.0,
            output_per_token: 4.40 / 1_000_000.0,
        }
    } else if lower.starts_with("o1-mini") {
        OpenAiRates {
            input_per_token: 3.00 / 1_000_000.0,
            output_per_token: 12.00 / 1_000_000.0,
        }
    } else if lower.starts_with("o1") {
        OpenAiRates {
            input_per_token: 15.00 / 1_000_000.0,
            output_per_token: 60.00 / 1_000_000.0,
        }
    } else {
        OpenAiRates {
            input_per_token: 0.15 / 1_000_000.0,
            output_per_token: 0.60 / 1_000_000.0,
        }
    }
}

fn build_chat_body(messages: &[Message], options: &CompletionOptions) -> Value {
    let msgs: Vec<Value> = messages
        .iter()
        .map(|m| {
            json!({
                "role": match m.role {
                    Role::System => "system",
                    Role::User => "user",
                    Role::Assistant => "assistant",
                    // The agent conversation model is plain text (Message
                    // carries no `tool_call_id` or assistant `tool_calls`), so
                    // a tool result is sent as a user turn. OpenAI rejects a
                    // `tool`-role message that is not paired with a preceding
                    // assistant `tool_calls`; sending it as `user` (as the
                    // Anthropic provider also does) keeps the multi-turn agent
                    // loop working without that structured pairing.
                    Role::Tool => "user",
                },
                "content": m.content,
            })
        })
        .collect();
    let mut body = json!({
        "model": options.model,
        "messages": msgs,
    });
    if let Some(t) = options.temperature {
        body["temperature"] = json!(t);
    }
    if let Some(n) = options.max_tokens {
        body["max_tokens"] = json!(n);
    }
    if options.json_mode {
        body["response_format"] = json!({"type": "json_object"});
    }
    if !options.tools.is_empty() {
        body["tools"] = options
            .tools
            .iter()
            .map(|t| {
                json!({
                    "type": "function",
                    "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.input_schema,
                    },
                })
            })
            .collect();
        body["tool_choice"] = match &options.tool_choice {
            Some(ToolChoice::Tool(name)) => json!({
                "type": "function",
                "function": {"name": name},
            }),
            Some(ToolChoice::Required) => json!("required"),
            Some(ToolChoice::Auto) | None => json!("auto"),
        };
    }
    body
}

async fn handle_response(resp: reqwest::Response) -> Result<Value, ProviderError> {
    let status = resp.status();
    let retry_after = parse_retry_after(resp.headers().get("retry-after"));
    if status.is_success() {
        return resp
            .json::<Value>()
            .await
            .map_err(|e| ProviderError::MalformedResponse(e.to_string()));
    }
    let body = resp.text().await.unwrap_or_default();
    match status.as_u16() {
        401 | 403 => Err(ProviderError::AuthFailure(body)),
        429 => Err(ProviderError::RateLimit { retry_after }),
        400 | 404 | 422 => Err(ProviderError::InvalidRequest(body)),
        500..=599 => Err(ProviderError::ProviderInternal(body)),
        _ => Err(ProviderError::ProviderInternal(format!(
            "unexpected status {status}: {body}"
        ))),
    }
}

fn parse_retry_after(header: Option<&reqwest::header::HeaderValue>) -> Option<Duration> {
    header
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.parse::<u64>().ok())
        .map(Duration::from_secs)
}

#[derive(Deserialize)]
struct OpenAiUsage {
    #[serde(default)]
    prompt_tokens: u32,
    #[serde(default)]
    completion_tokens: u32,
}

fn parse_chat_response(value: Value, requested_model: &str) -> Result<Response, ProviderError> {
    let model = value
        .get("model")
        .and_then(|v| v.as_str())
        .unwrap_or(requested_model)
        .to_owned();
    let choice = value
        .get("choices")
        .and_then(|v| v.as_array())
        .and_then(|a| a.first())
        .ok_or_else(|| ProviderError::MalformedResponse("missing choices[0]".into()))?;
    let finish = choice
        .get("finish_reason")
        .and_then(|v| v.as_str())
        .unwrap_or("stop");
    let message = choice
        .get("message")
        .ok_or_else(|| ProviderError::MalformedResponse("missing message".into()))?;
    let usage_obj: OpenAiUsage = serde_json::from_value(
        value
            .get("usage")
            .cloned()
            .unwrap_or_else(|| Value::Object(Default::default())),
    )
    .map_err(|e| ProviderError::MalformedResponse(format!("usage parse: {e}")))?;
    let usage = Usage {
        input_tokens: usage_obj.prompt_tokens,
        output_tokens: usage_obj.completion_tokens,
        cache_read_tokens: None,
        cache_write_tokens: None,
    };
    let content = extract_content(message)?;
    Ok(Response {
        content,
        finish_reason: map_finish_reason(finish),
        usage,
        model,
    })
}

fn extract_content(message: &Value) -> Result<ResponseContent, ProviderError> {
    if let Some(calls) = message.get("tool_calls").and_then(|v| v.as_array()) {
        let parsed: Result<Vec<ToolCall>, ProviderError> = calls
            .iter()
            .map(|c| {
                let name = c
                    .pointer("/function/name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_owned();
                let raw_args = c
                    .pointer("/function/arguments")
                    .and_then(|v| v.as_str())
                    .unwrap_or("{}");
                let arguments: Value = serde_json::from_str(raw_args)
                    .map_err(|e| ProviderError::MalformedResponse(format!("tool args: {e}")))?;
                Ok(ToolCall { name, arguments })
            })
            .collect();
        let calls = parsed?;
        if !calls.is_empty() {
            return Ok(ResponseContent::ToolUse { calls });
        }
    }
    let text = message
        .get("content")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_owned();
    Ok(ResponseContent::Text { text })
}

fn map_finish_reason(s: &str) -> FinishReason {
    match s {
        "stop" => FinishReason::Stop,
        "length" => FinishReason::MaxTokens,
        "tool_calls" | "function_call" => FinishReason::ToolUse,
        _ => FinishReason::Other,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eval::provider::ToolSpec;
    use serde_json::json;
    use wiremock::matchers::{body_partial_json, header, method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    fn provider_for(server: &MockServer) -> OpenAiProvider {
        OpenAiProvider::new(OpenAiConfig {
            api_key: "test-key".into(),
            organization: None,
            base_url: Some(server.uri()),
            request_timeout: Duration::from_secs(5),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap()
    }

    #[tokio::test]
    async fn happy_path_chat_completion_returns_text() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .and(header("authorization", "Bearer test-key"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "gpt-4o-mini",
                "choices": [{
                    "message": {"role": "assistant", "content": "hello"},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 5},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let messages = vec![Message {
            role: Role::User,
            content: "hi".into(),
        }];
        let options = CompletionOptions {
            model: "gpt-4o-mini".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&messages, &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "hello"),
            other => panic!("expected text, got {other:?}"),
        }
        assert_eq!(response.usage.output_tokens, 5);
    }

    #[tokio::test]
    async fn json_mode_sets_response_format() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .and(body_partial_json(json!({
                "response_format": {"type": "json_object"}
            })))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "gpt-4o",
                "choices": [{
                    "message": {"role": "assistant", "content": "{\"ok\": true}"},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 5, "completion_tokens": 5},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gpt-4o".into(),
            json_mode: true,
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert!(text.contains("\"ok\"")),
            other => panic!("expected text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn count_tokens_uses_o200k_for_gpt_4o() {
        let provider = OpenAiProvider::new(OpenAiConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            ..Default::default()
        })
        .unwrap();
        let count = provider
            .count_tokens("hello world", "gpt-4o")
            .await
            .unwrap();
        // Known fixture: "hello world" is 2 tokens under o200k_base.
        assert_eq!(count, 2);
    }

    #[tokio::test]
    async fn count_tokens_uses_cl100k_for_gpt_4() {
        let provider = OpenAiProvider::new(OpenAiConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            ..Default::default()
        })
        .unwrap();
        let count = provider.count_tokens("hello world", "gpt-4").await.unwrap();
        // cl100k splits "hello world" the same way: 2 tokens.
        assert_eq!(count, 2);
    }

    #[tokio::test]
    async fn rate_limit_retries_then_succeeds() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(
                ResponseTemplate::new(429)
                    .insert_header("retry-after", "0")
                    .set_body_string("slow down"),
            )
            .up_to_n_times(2)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "gpt-4o-mini",
                "choices": [{
                    "message": {"role": "assistant", "content": "ok"},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 1, "completion_tokens": 1},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gpt-4o-mini".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "ok"),
            other => panic!("expected text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn azure_base_url_smoke() {
        // Smoke: a custom base URL works the same way as the OpenAI one;
        // wiremock stands in for the Azure resource so we exercise the
        // header + path the Azure deployment expects.
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "gpt-4o-azure",
                "choices": [{
                    "message": {"role": "assistant", "content": "azure-ok"},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 1, "completion_tokens": 1},
            })))
            .mount(&server)
            .await;
        let provider = OpenAiProvider::new(OpenAiConfig {
            api_key: "azure-key".into(),
            organization: Some("org-1".into()),
            base_url: Some(server.uri()),
            request_timeout: Duration::from_secs(5),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap();
        let options = CompletionOptions {
            model: "gpt-4o".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "azure-ok"),
            other => panic!("expected text, got {other:?}"),
        }
    }

    #[test]
    fn tool_role_messages_are_sent_as_user_turns() {
        // OpenAI rejects a `tool`-role message that is not paired with a
        // preceding assistant `tool_calls`. The plain-text agent Message model
        // carries no such pairing, so a tool result must go out as a user turn
        // (matching the Anthropic provider), keeping the multi-turn loop valid.
        let messages = vec![
            Message {
                role: Role::User,
                content: "what is the weather?".into(),
            },
            Message {
                role: Role::Tool,
                content: "{\"city\":\"Sacramento\",\"temp\":72}".into(),
            },
        ];
        let body = build_chat_body(&messages, &CompletionOptions::default());
        let roles: Vec<&str> = body["messages"]
            .as_array()
            .unwrap()
            .iter()
            .map(|m| m["role"].as_str().unwrap())
            .collect();
        assert_eq!(
            roles,
            vec!["user", "user"],
            "tool result maps to a user turn"
        );
    }

    #[test]
    fn forced_tool_call_round_trips_through_body_and_response() {
        // Request side: a forced tool choice serializes the tools array
        // and pins tool_choice to the named function.
        let options = CompletionOptions {
            model: "gpt-4o".into(),
            tools: vec![ToolSpec {
                name: "report".into(),
                description: "Return a verdict".into(),
                input_schema: json!({"type": "object"}),
            }],
            tool_choice: Some(ToolChoice::Tool("report".into())),
            ..Default::default()
        };
        let body = build_chat_body(&[], &options);
        assert_eq!(
            body.pointer("/tools/0/function/name")
                .and_then(|v| v.as_str()),
            Some("report"),
            "tools must be serialized into the request body"
        );
        assert_eq!(
            body.pointer("/tool_choice/function/name")
                .and_then(|v| v.as_str()),
            Some("report")
        );

        // Response side: a canned tool_calls response parses into ToolUse.
        let canned = json!({
            "model": "gpt-4o",
            "choices": [{
                "message": {
                    "role": "assistant",
                    "tool_calls": [{
                        "type": "function",
                        "function": {
                            "name": "report",
                            "arguments": "{\"verdict\":\"pass\"}",
                        },
                    }],
                },
                "finish_reason": "tool_calls",
            }],
            "usage": {"prompt_tokens": 5, "completion_tokens": 3},
        });
        let response = parse_chat_response(canned, "gpt-4o").unwrap();
        match response.content {
            ResponseContent::ToolUse { calls } => {
                assert_eq!(calls.len(), 1);
                assert_eq!(calls[0].name, "report");
                assert_eq!(calls[0].arguments, json!({"verdict": "pass"}));
            }
            other => panic!("expected tool use, got {other:?}"),
        }
        assert_eq!(response.finish_reason, FinishReason::ToolUse);
    }

    #[tokio::test]
    async fn embed_parses_vector_from_data_array() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/embeddings"))
            .and(header("authorization", "Bearer test-key"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "data": [{ "embedding": [0.1, 0.2, 0.3] }],
                "model": "text-embedding-3-small",
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let v = provider
            .embed("hello", "text-embedding-3-small")
            .await
            .unwrap();
        assert_eq!(v.len(), 3);
        assert!((v[1] - 0.2).abs() < 1e-6);
    }

    #[test]
    fn estimate_cost_uses_gpt4o_mini_rates_for_default_model() {
        let provider = OpenAiProvider::new(OpenAiConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            ..Default::default()
        })
        .unwrap();
        let usage = Usage {
            input_tokens: 1_000_000,
            output_tokens: 1_000_000,
            ..Default::default()
        };
        let cost = provider.estimate_cost(&usage, "gpt-4o-mini");
        // $0.15 + $0.60 = $0.75 per pair of one million tokens.
        assert!((cost.total() - 0.75).abs() < 1e-6);
    }
}
