//! Anthropic provider impl.
//!
//! Wraps the Messages API (`/v1/messages`) and the token-count endpoint
//! (`/v1/messages/count_tokens`). Tool use is the route to structured
//! output because Anthropic does not expose a JSON-mode flag; the runner
//! forces a `tool_choice` and reads the JSON arguments back from the
//! response. Pricing follows the published per-million rate cards; values
//! are kept in the table at the bottom of the file so a price change is
//! a one-line edit.

use std::time::Duration;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
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::retry::run_with_retry;

const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
const ANTHROPIC_VERSION: &str = "2023-06-01";

/// Anthropic provider client.
///
/// Construct via [`AnthropicProvider::new`]. Holds a `reqwest::Client`
/// that reuses connections across calls; provider impls are designed to
/// be long-lived (one per juror).
pub struct AnthropicProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
    retry: RetryPolicy,
    request_timeout: Duration,
}

/// Config bag passed to [`AnthropicProvider::new`].
///
/// `api_key` is resolved by the loader from `${env:ANTHROPIC_API_KEY}`;
/// the loader fails fast when the env var is missing so we never get
/// halfway through a CI run before noticing.
#[derive(Debug, Clone)]
pub struct AnthropicConfig {
    /// API key, already resolved from the env.
    pub api_key: String,
    /// Override the base URL (useful for Bedrock-fronted gateways and
    /// tests against `wiremock`).
    pub base_url: Option<String>,
    /// HTTP request timeout. Default 30s for remote calls.
    pub request_timeout: Duration,
    /// HTTP / HTTPS proxy policy. Empty by default; reqwest
    /// keeps reading `HTTPS_PROXY` from the environment in that case.
    pub proxy: crate::network::ProxyConfig,
}

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

impl AnthropicProvider {
    /// Build a new provider client. Returns [`ProviderError::Transport`]
    /// when the underlying `reqwest::Client` cannot be constructed (TLS
    /// initialization failure).
    pub fn new(config: AnthropicConfig) -> 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,
            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 AnthropicProvider {
    fn name(&self) -> ProviderId {
        ProviderId::Anthropic
    }

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

    async fn chat_completion(
        &self,
        messages: &[Message],
        options: &CompletionOptions,
    ) -> Result<Response, ProviderError> {
        let body = build_messages_body(messages, options);
        let url = format!("{}/v1/messages", self.base_url);
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .header("x-api-key", &self.api_key)
                .header("anthropic-version", ANTHROPIC_VERSION)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_messages_response(response_json, &options.model)
    }

    async fn count_tokens(&self, text: &str, model: &str) -> Result<u32, ProviderError> {
        let url = format!("{}/v1/messages/count_tokens", self.base_url);
        let body = json!({
            "model": model,
            "messages": [{"role": "user", "content": text}],
        });
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .header("x-api-key", &self.api_key)
                .header("anthropic-version", ANTHROPIC_VERSION)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        response_json
            .get("input_tokens")
            .and_then(|v| v.as_u64())
            .map(|n| n as u32)
            .ok_or_else(|| {
                ProviderError::MalformedResponse("count_tokens missing input_tokens".into())
            })
    }

    async fn embed(&self, _text: &str, _model: &str) -> Result<Vec<f32>, ProviderError> {
        // Anthropic ships no embeddings endpoint; they point users at Voyage
        // AI instead. Return a clear unsupported error rather than faking a
        // vector so the `similar` matcher fails loudly on a bad provider pick.
        Err(ProviderError::FeatureUnsupported {
            provider: ProviderId::Anthropic.as_str(),
            feature: "embeddings",
        })
    }

    fn estimate_cost(&self, usage: &Usage, model: &str) -> Cost {
        let rates = anthropic_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: usage
                .cache_read_tokens
                .map(|n| rates.cache_read_per_token * f64::from(n)),
            cache_write_dollars: usage
                .cache_write_tokens
                .map(|n| rates.cache_write_per_token * f64::from(n)),
        }
    }
}

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

/// Published pricing (USD per token) for the common Claude models. Add
/// a row when a new model name lands; unknown models fall back to the
/// Sonnet rates, which is the most conservative default.
fn anthropic_rates(model: &str) -> AnthropicRates {
    let lower = model.to_ascii_lowercase();
    if lower.contains("haiku") {
        // Claude 3.5 Haiku: $0.80 / $4 per million tokens.
        AnthropicRates {
            input_per_token: 0.80 / 1_000_000.0,
            output_per_token: 4.00 / 1_000_000.0,
            cache_read_per_token: 0.08 / 1_000_000.0,
            cache_write_per_token: 1.00 / 1_000_000.0,
        }
    } else if lower.contains("opus") {
        // Claude 3 Opus: $15 / $75 per million tokens.
        AnthropicRates {
            input_per_token: 15.00 / 1_000_000.0,
            output_per_token: 75.00 / 1_000_000.0,
            cache_read_per_token: 1.50 / 1_000_000.0,
            cache_write_per_token: 18.75 / 1_000_000.0,
        }
    } else {
        // Sonnet (default): $3 / $15 per million tokens.
        AnthropicRates {
            input_per_token: 3.00 / 1_000_000.0,
            output_per_token: 15.00 / 1_000_000.0,
            cache_read_per_token: 0.30 / 1_000_000.0,
            cache_write_per_token: 3.75 / 1_000_000.0,
        }
    }
}

#[derive(Serialize)]
struct AnthropicMessage<'a> {
    role: &'a str,
    content: &'a str,
}

fn build_messages_body(messages: &[Message], options: &CompletionOptions) -> Value {
    let (system, chat): (Vec<&Message>, Vec<&Message>) = messages
        .iter()
        .partition(|m| matches!(m.role, Role::System));
    let system_prompt = system
        .iter()
        .map(|m| m.content.as_str())
        .collect::<Vec<_>>()
        .join("\n\n");
    let chat_msgs: Vec<AnthropicMessage> = chat
        .iter()
        .map(|m| AnthropicMessage {
            role: match m.role {
                Role::User => "user",
                Role::Assistant => "assistant",
                Role::Tool => "user",
                Role::System => "user",
            },
            content: &m.content,
        })
        .collect();
    let mut body = json!({
        "model": options.model,
        "messages": chat_msgs,
        "max_tokens": options.max_tokens.unwrap_or(1024),
    });
    if !system_prompt.is_empty() {
        body["system"] = Value::String(system_prompt);
    }
    if let Some(t) = options.temperature {
        body["temperature"] = json!(t);
    }
    if !options.tools.is_empty() {
        body["tools"] = options
            .tools
            .iter()
            .map(|t| {
                json!({
                    "name": t.name,
                    "description": t.description,
                    "input_schema": t.input_schema,
                })
            })
            .collect();
        body["tool_choice"] = match &options.tool_choice {
            Some(ToolChoice::Tool(name)) => json!({"type": "tool", "name": name}),
            Some(ToolChoice::Required) => json!({"type": "any"}),
            Some(ToolChoice::Auto) | None => json!({"type": "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 AnthropicUsage {
    #[serde(default)]
    input_tokens: u32,
    #[serde(default)]
    output_tokens: u32,
    #[serde(default)]
    cache_read_input_tokens: Option<u32>,
    #[serde(default)]
    cache_creation_input_tokens: Option<u32>,
}

fn parse_messages_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 stop_reason = value
        .get("stop_reason")
        .and_then(|v| v.as_str())
        .unwrap_or("end_turn");
    let usage_obj: AnthropicUsage = 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.input_tokens,
        output_tokens: usage_obj.output_tokens,
        cache_read_tokens: usage_obj.cache_read_input_tokens,
        cache_write_tokens: usage_obj.cache_creation_input_tokens,
    };
    let content = extract_content(&value)?;
    Ok(Response {
        content,
        finish_reason: map_stop_reason(stop_reason),
        usage,
        model,
    })
}

fn extract_content(value: &Value) -> Result<ResponseContent, ProviderError> {
    let blocks = value
        .get("content")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ProviderError::MalformedResponse("missing content array".into()))?;
    let mut text_parts: Vec<String> = Vec::new();
    let mut tool_calls: Vec<ToolCall> = Vec::new();
    for block in blocks {
        match block.get("type").and_then(|v| v.as_str()) {
            Some("text") => {
                if let Some(t) = block.get("text").and_then(|v| v.as_str()) {
                    text_parts.push(t.to_owned());
                }
            }
            Some("tool_use") => {
                let name = block
                    .get("name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_owned();
                let arguments = block
                    .get("input")
                    .cloned()
                    .unwrap_or(Value::Object(Default::default()));
                tool_calls.push(ToolCall { name, arguments });
            }
            _ => {}
        }
    }
    if !tool_calls.is_empty() {
        return Ok(ResponseContent::ToolUse { calls: tool_calls });
    }
    Ok(ResponseContent::Text {
        text: text_parts.join(""),
    })
}

fn map_stop_reason(s: &str) -> FinishReason {
    match s {
        "end_turn" | "stop_sequence" => FinishReason::Stop,
        "max_tokens" => FinishReason::MaxTokens,
        "tool_use" => FinishReason::ToolUse,
        _ => FinishReason::Other,
    }
}

#[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) -> AnthropicProvider {
        AnthropicProvider::new(AnthropicConfig {
            api_key: "test-key".into(),
            base_url: Some(server.uri()),
            request_timeout: Duration::from_secs(5),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap()
    }

    #[tokio::test]
    async fn happy_path_messages_call_returns_text() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/messages"))
            .and(header("x-api-key", "test-key"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "msg_01",
                "model": "claude-3-5-sonnet-20241022",
                "content": [{"type": "text", "text": "hi there"}],
                "stop_reason": "end_turn",
                "usage": {"input_tokens": 12, "output_tokens": 5},
            })))
            .mount(&server)
            .await;

        let provider = provider_for(&server);
        let messages = vec![Message {
            role: Role::User,
            content: "say hi".into(),
        }];
        let options = CompletionOptions {
            model: "claude-3-5-sonnet-20241022".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&messages, &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "hi there"),
            other => panic!("expected text content, got {other:?}"),
        }
        assert_eq!(response.usage.input_tokens, 12);
        assert_eq!(response.finish_reason, FinishReason::Stop);
    }

    #[tokio::test]
    async fn tool_use_response_is_parsed_as_tool_calls() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/messages"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "msg_02",
                "model": "claude-3-5-sonnet",
                "content": [
                    {"type": "tool_use", "name": "report", "input": {"verdict": "pass"}},
                ],
                "stop_reason": "tool_use",
                "usage": {"input_tokens": 30, "output_tokens": 9},
            })))
            .mount(&server)
            .await;

        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "claude-3-5-sonnet".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 response = provider.chat_completion(&[], &options).await.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);
    }

    #[test]
    fn forced_tool_call_serializes_tools_and_tool_choice() {
        // Anthropic advertises tool_use; a forced tool choice serializes
        // the tools array with input_schema and pins tool_choice by name.
        let options = CompletionOptions {
            model: "claude-3-5-sonnet".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_messages_body(&[], &options);
        assert_eq!(
            body.pointer("/tools/0/name").and_then(|v| v.as_str()),
            Some("report"),
            "tools must be serialized into the request body"
        );
        assert!(body.pointer("/tools/0/input_schema").is_some());
        assert_eq!(
            body.pointer("/tool_choice/name").and_then(|v| v.as_str()),
            Some("report")
        );
    }

    #[tokio::test]
    async fn count_tokens_uses_dedicated_endpoint() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/messages/count_tokens"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({"input_tokens": 17})))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let count = provider
            .count_tokens("hello world", "claude-3-5-sonnet")
            .await
            .unwrap();
        assert_eq!(count, 17);
    }

    #[tokio::test]
    async fn rate_limit_with_retry_after_eventually_succeeds() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/messages"))
            .respond_with(
                ResponseTemplate::new(429)
                    .insert_header("retry-after", "0")
                    .set_body_string("rate limited"),
            )
            .up_to_n_times(2)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(path("/v1/messages"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "claude-3-5-sonnet",
                "content": [{"type": "text", "text": "ok"}],
                "stop_reason": "end_turn",
                "usage": {"input_tokens": 1, "output_tokens": 1},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "claude-3-5-sonnet".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 auth_failure_does_not_retry() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/messages"))
            .respond_with(ResponseTemplate::new(401).set_body_string("bad key"))
            .expect(1)
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "claude-3-5-sonnet".into(),
            ..Default::default()
        };
        let err = provider.chat_completion(&[], &options).await.unwrap_err();
        assert!(matches!(err, ProviderError::AuthFailure(_)));
    }

    #[test]
    fn estimate_cost_uses_sonnet_rates_for_default_model() {
        let provider = AnthropicProvider::new(AnthropicConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            request_timeout: Duration::from_secs(1),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap();
        let usage = Usage {
            input_tokens: 1_000_000,
            output_tokens: 1_000_000,
            ..Default::default()
        };
        let cost = provider.estimate_cost(&usage, "claude-3-5-sonnet");
        // 1M input @ $3 + 1M output @ $15 = $18.
        assert!((cost.total() - 18.0).abs() < 1e-6);
    }

    #[tokio::test]
    async fn embed_reports_feature_unsupported() {
        // Anthropic has no embeddings endpoint; `embed` must error rather than
        // touch the network, so no mock server is needed.
        let provider = AnthropicProvider::new(AnthropicConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            request_timeout: Duration::from_secs(1),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap();
        let err = provider
            .embed("hello", "text-embedding-3-small")
            .await
            .unwrap_err();
        assert!(matches!(err, ProviderError::FeatureUnsupported { .. }));
    }
}
