//! Mistral provider impl.
//!
//! The Mistral chat API is close enough to OpenAI's that the request and
//! response shapes round-trip with minimal massaging. We still keep an
//! independent implementation so a future divergence does not force a
//! shared OpenAI/Mistral helper to grow conditional branches.

use std::time::Duration;

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

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.mistral.ai";

/// Mistral provider client.
pub struct MistralProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
    retry: RetryPolicy,
    request_timeout: Duration,
}

/// Config bag for [`MistralProvider::new`].
#[derive(Debug, Clone)]
pub struct MistralConfig {
    /// Resolved API key.
    pub api_key: String,
    /// Override base URL (wiremock or self-hosted Mistral gateway).
    pub base_url: Option<String>,
    /// HTTP timeout for body reads. Default 120s (slow Mistral models).
    pub request_timeout: Duration,
    /// HTTP connect timeout. Default 30s.
    pub connect_timeout: Duration,
    /// HTTP / HTTPS proxy policy.
    pub proxy: crate::network::ProxyConfig,
}

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

impl MistralProvider {
    /// Build a new client.
    pub fn new(config: MistralConfig) -> Result<Self, ProviderError> {
        let builder = reqwest::Client::builder()
            .timeout(config.request_timeout)
            .connect_timeout(config.connect_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,
            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 MistralProvider {
    fn name(&self) -> ProviderId {
        ProviderId::Mistral
    }

    fn supported_capabilities(&self) -> Capabilities {
        Capabilities {
            json_mode: true,
            tool_use: true,
            streaming: true,
            structured_output: true,
            token_count_exact: false,
        }
    }

    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 resp = self
                .client
                .post(&url)
                .bearer_auth(&self.api_key)
                .json(&body)
                .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> {
        // Mistral has no public tokenizer endpoint; fall back to a
        // cl100k_base estimate. Documented as approximate in ADR-0020.
        let bpe = tiktoken_rs::cl100k_base()
            .map_err(|e| ProviderError::Transport(format!("tiktoken: {e}")))?;
        Ok(bpe.encode_with_special_tokens(text).len() 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 resp = self
                .client
                .post(&url)
                .bearer_auth(&self.api_key)
                .json(&body)
                .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 = mistral_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,
        }
    }
}

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

fn mistral_rates(model: &str) -> MistralRates {
    let lower = model.to_ascii_lowercase();
    if lower.contains("small") {
        MistralRates {
            input_per_token: 0.20 / 1_000_000.0,
            output_per_token: 0.60 / 1_000_000.0,
        }
    } else if lower.contains("nemo") {
        MistralRates {
            input_per_token: 0.15 / 1_000_000.0,
            output_per_token: 0.15 / 1_000_000.0,
        }
    } else {
        // Mistral-large (default).
        MistralRates {
            input_per_token: 2.00 / 1_000_000.0,
            output_per_token: 6.00 / 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",
                    Role::Tool => "tool",
                },
                "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() {
        // Mistral's tools array is OpenAI-shaped, so this mirrors the
        // OpenAI provider field for field.
        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)),
        404 => {
            // Mistral returns 404 with a body that mentions "model" when
            // the requested model is unknown to the gateway. Map that to
            // an InvalidRequest so the runner does not retry.
            Err(ProviderError::InvalidRequest(format!(
                "model not found: {body}"
            )))
        }
        429 => Err(ProviderError::RateLimit { retry_after }),
        400 | 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 MistralUsage {
    #[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: MistralUsage = 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 content = extract_content(message)?;
    Ok(Response {
        content,
        finish_reason: match finish {
            "stop" => FinishReason::Stop,
            "length" => FinishReason::MaxTokens,
            "tool_calls" => FinishReason::ToolUse,
            _ => FinishReason::Other,
        },
        usage: Usage {
            input_tokens: usage_obj.prompt_tokens,
            output_tokens: usage_obj.completion_tokens,
            cache_read_tokens: None,
            cache_write_tokens: None,
        },
        model,
    })
}

fn extract_content(message: &Value) -> Result<ResponseContent, ProviderError> {
    // Mistral mirrors the OpenAI tool_calls shape: each call carries a
    // function name and a JSON-encoded argument string.
    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 })
}

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

    fn provider_for(server: &MockServer) -> MistralProvider {
        MistralProvider::new(MistralConfig {
            api_key: "test-key".into(),
            base_url: Some(server.uri()),
            request_timeout: Duration::from_secs(5),
            connect_timeout: Duration::from_secs(1),
            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": "mistral-large-latest",
                "choices": [{
                    "message": {"role": "assistant", "content": "bonjour"},
                    "finish_reason": "stop",
                }],
                "usage": {"prompt_tokens": 3, "completion_tokens": 2},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let messages = vec![Message {
            role: Role::User,
            content: "salut".into(),
        }];
        let options = CompletionOptions {
            model: "mistral-large-latest".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&messages, &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "bonjour"),
            other => panic!("expected text, got {other:?}"),
        }
        assert_eq!(response.usage.input_tokens, 3);
    }

    #[test]
    fn forced_tool_call_round_trips_through_body_and_response() {
        // Request side: Mistral advertises tool_use, so a forced tool
        // choice must serialize the OpenAI-shaped tools array rather than
        // silently dropping it.
        let options = CompletionOptions {
            model: "mistral-large-latest".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": "mistral-large-latest",
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": "",
                    "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, "mistral-large-latest").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 model_not_found_maps_to_invalid_request() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(ResponseTemplate::new(404).set_body_string("Model nope not found"))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "nope".into(),
            ..Default::default()
        };
        let err = provider.chat_completion(&[], &options).await.unwrap_err();
        assert!(matches!(err, ProviderError::InvalidRequest(_)));
    }
}
