//! `LlmProvider` trait and its supporting value types (ADR-0020).
//!
//! Why this module exists: the `llm-judge` matcher and the
//! model-compatibility feature both invoke language models. If
//! provider-specific request shapes leak into the matcher, every new
//! provider adds a code path. The trait below pins one request shape, one
//! response shape, and one error taxonomy. Provider modules implement the
//! trait and map their native quirks at the boundary.

use std::time::Duration;

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

use super::error::ProviderError;

/// Stable identifier for an LLM provider.
///
/// Surfaced in config, reporter output, and logs. Kept kebab-case so the
/// JSON reporter renders the same string a YAML config writes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
#[non_exhaustive]
pub enum ProviderId {
    /// Anthropic Messages API.
    Anthropic,
    /// OpenAI Chat Completions API (also covers Azure OpenAI via base URL).
    Openai,
    /// Mistral chat completions API.
    Mistral,
    /// Local Ollama HTTP API.
    Ollama,
    /// Google Gemini `generateContent` API.
    Google,
}

impl ProviderId {
    /// Return the kebab-case wire string for this provider.
    pub fn as_str(self) -> &'static str {
        match self {
            ProviderId::Anthropic => "anthropic",
            ProviderId::Openai => "openai",
            ProviderId::Mistral => "mistral",
            ProviderId::Ollama => "ollama",
            ProviderId::Google => "google",
        }
    }
}

/// Capability bitflags advertised by a provider+model pair.
///
/// The runner reads these before dispatching work to a juror. A juror that
/// requests `tool_use` against a provider whose impl reports `tool_use:
/// false` is rejected at load time with a clean error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Capabilities {
    /// Provider supports a JSON-mode response format.
    pub json_mode: bool,
    /// Provider supports tool/function calling.
    pub tool_use: bool,
    /// Provider supports server-sent token streaming.
    pub streaming: bool,
    /// Provider returns structured output (either via JSON mode or by
    /// forcing a tool call).
    pub structured_output: bool,
    /// Provider can count tokens exactly (versus a tiktoken estimate).
    pub token_count_exact: bool,
}

impl Capabilities {
    /// Conservative defaults: nothing supported. Provider impls override
    /// in their constructor so the bitflags are an explicit opt-in.
    pub const fn none() -> Self {
        Self {
            json_mode: false,
            tool_use: false,
            streaming: false,
            structured_output: false,
            token_count_exact: false,
        }
    }
}

/// Conversation role used in a [`Message`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System prompt; sets persona and rules.
    System,
    /// User turn.
    User,
    /// Assistant turn (replayed history).
    Assistant,
    /// Tool result fed back to the model.
    Tool,
}

/// One message in a chat-style request.
///
/// Tool use and structured output ride on top of `content` and the
/// `options.tools` slice. The trait deliberately keeps the message shape
/// flat; provider impls translate to the native request body.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Message {
    /// Role of the speaker.
    pub role: Role,
    /// Plain-text content. Tool-call requests pack the JSON tool input
    /// here and let the provider impl extract it.
    pub content: String,
}

/// Tool descriptor passed alongside a completion request.
///
/// Tool-use forcing (Anthropic's "force a specific tool" + OpenAI's
/// `tool_choice`) goes through [`CompletionOptions::tool_choice`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolSpec {
    /// Tool name. Must be a stable identifier the caller can match on.
    pub name: String,
    /// Free-form description shown to the model.
    pub description: String,
    /// JSON Schema for the tool's input arguments.
    pub input_schema: Value,
}

/// Selection policy for tool calls.
///
/// `Auto` lets the model decide; `Required` forces some tool call; `Tool`
/// forces a specific tool by name (the Anthropic-and-OpenAI pattern that
/// also gives Anthropic JSON-mode parity).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolChoice {
    /// Model decides whether to call a tool.
    Auto,
    /// Model must call some tool.
    Required,
    /// Model must call the named tool.
    Tool(String),
}

/// Options applied to a single completion call.
#[derive(Debug, Clone, Default)]
pub struct CompletionOptions {
    /// Model identifier (provider-specific; `claude-3-5-sonnet`,
    /// `gpt-4o`, `mistral-large-latest`, etc.).
    pub model: String,
    /// Sampling temperature in `[0.0, 2.0]`; `None` defers to provider
    /// default.
    pub temperature: Option<f32>,
    /// Hard cap on output tokens.
    pub max_tokens: Option<u32>,
    /// Request structured JSON output (provider-specific flag).
    pub json_mode: bool,
    /// Tool descriptors offered to the model.
    pub tools: Vec<ToolSpec>,
    /// Tool selection policy, when `tools` is non-empty.
    pub tool_choice: Option<ToolChoice>,
}

/// Tool call extracted from a model response.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCall {
    /// Name of the tool the model decided to call.
    pub name: String,
    /// JSON arguments the model passed to the tool.
    pub arguments: Value,
}

/// Response content normalized across providers.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "kebab-case")]
pub enum ResponseContent {
    /// Plain text reply.
    Text {
        /// The text body returned by the model.
        text: String,
    },
    /// One or more tool calls produced by the model.
    ToolUse {
        /// Tool calls in the order the model emitted them.
        calls: Vec<ToolCall>,
    },
}

/// Why the model stopped generating.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
#[non_exhaustive]
pub enum FinishReason {
    /// Model hit a natural stop point.
    Stop,
    /// Output hit `max_tokens`.
    MaxTokens,
    /// Model decided to call a tool.
    ToolUse,
    /// Provider returned a finish reason we did not recognize; the
    /// original string lives in the response payload for debugging.
    Other,
}

/// Token-usage breakdown returned alongside every completion.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    /// Input (prompt) tokens billed by the provider.
    pub input_tokens: u32,
    /// Output (completion) tokens billed by the provider.
    pub output_tokens: u32,
    /// Tokens served from a provider-side prompt cache, when supported
    /// (Anthropic; other providers leave this `None`).
    pub cache_read_tokens: Option<u32>,
    /// Tokens written into a provider-side prompt cache, when supported.
    pub cache_write_tokens: Option<u32>,
}

/// Dollar cost breakdown derived from a [`Usage`] record.
///
/// Kept structured so the reporter can render "of the $0.12 you spent,
/// $0.04 was prompt-cache writes" without re-deriving the split.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct Cost {
    /// Cost attributable to input tokens.
    pub input_dollars: f64,
    /// Cost attributable to output tokens.
    pub output_dollars: f64,
    /// Cost attributable to prompt-cache reads, when supported.
    pub cache_read_dollars: Option<f64>,
    /// Cost attributable to prompt-cache writes, when supported.
    pub cache_write_dollars: Option<f64>,
}

impl Cost {
    /// Total cost across every line item.
    pub fn total(&self) -> f64 {
        self.input_dollars
            + self.output_dollars
            + self.cache_read_dollars.unwrap_or(0.0)
            + self.cache_write_dollars.unwrap_or(0.0)
    }
}

/// Response returned by [`LlmProvider::chat_completion`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Response {
    /// Content the model returned (text or tool calls).
    pub content: ResponseContent,
    /// Why generation stopped.
    pub finish_reason: FinishReason,
    /// Token usage echoed by the provider.
    pub usage: Usage,
    /// Model identifier the provider actually served; useful when the
    /// caller asked for an alias that resolved to a specific snapshot.
    pub model: String,
}

/// Retry policy applied to a single completion call.
///
/// Provider impls keep one instance per client. Defaults match the
/// per-provider direction/321/323: retry on rate limits,
/// timeouts, and provider-internal errors; never retry on auth or
/// malformed-request failures. The exact backoff schedule lives on each
/// provider's request loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RetryPolicy {
    /// Maximum retry attempts after the initial call.
    pub max_retries: u32,
    /// Base delay between retries; the call site applies exponential
    /// backoff against this base (50ms, 100ms, 200ms, ...).
    pub base_delay: Duration,
    /// Hard cap across all retries; if exceeded, return the last error.
    pub total_budget: Duration,
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_retries: 4,
            base_delay: Duration::from_millis(50),
            total_budget: Duration::from_secs(60),
        }
    }
}

impl RetryPolicy {
    /// Disable retries entirely. Used by the Ollama provider because
    /// local connection failures are not transient.
    pub fn none() -> Self {
        Self {
            max_retries: 0,
            base_delay: Duration::from_millis(0),
            total_budget: Duration::from_millis(0),
        }
    }

    /// Exponential delay for attempt `n` (zero-based). Returns
    /// `base_delay * 2^n`, saturating at u64::MAX milliseconds so a
    /// pathological policy cannot overflow.
    pub fn delay_for(&self, n: u32) -> Duration {
        let ms = self.base_delay.as_millis() as u64;
        let multiplier = 1u64.checked_shl(n).unwrap_or(u64::MAX);
        Duration::from_millis(ms.saturating_mul(multiplier))
    }
}

/// Map a `reqwest` transport failure onto the shared [`ProviderError`]
/// taxonomy.
///
/// Every HTTP-backed provider funnels send errors through this one
/// function so the timeout-versus-connect-versus-transport split cannot
/// drift between provider files. `configured_timeout` is the
/// provider's own request timeout, reported verbatim on a timeout so the
/// error message matches what the client was actually built with rather
/// than a hardcoded constant.
pub(crate) fn map_reqwest_error(
    err: reqwest::Error,
    configured_timeout: Duration,
) -> ProviderError {
    if err.is_timeout() {
        return ProviderError::Timeout(configured_timeout);
    }
    if err.is_connect() {
        // Keep connect failures tagged so a "connection refused" against a
        // local Ollama reads differently from a mid-flight TLS drop.
        return ProviderError::Transport(format!("connect: {err}"));
    }
    ProviderError::Transport(err.to_string())
}

/// Provider-side LLM client abstraction.
///
/// Implementations live under [`super::providers`]. The runner holds a
/// `Box<dyn LlmProvider>` per juror; nothing above that line knows which
/// provider is active. See ADR-0020 for the full rationale.
#[async_trait]
pub trait LlmProvider: Send + Sync {
    /// Stable identifier for this provider, used in config and logs.
    fn name(&self) -> ProviderId;

    /// Which optional features this provider+model supports. The runner
    /// reads this before deciding which juror gets which work item.
    fn supported_capabilities(&self) -> Capabilities;

    /// Single-shot completion. Most matcher calls land here.
    async fn chat_completion(
        &self,
        messages: &[Message],
        options: &CompletionOptions,
    ) -> Result<Response, ProviderError>;

    /// Count tokens for budget gating and prompt-fit checks. Providers
    /// that lack a tokenizer endpoint fall back to a tiktoken estimate
    /// and surface `Capabilities::token_count_exact = false`.
    async fn count_tokens(&self, text: &str, model: &str) -> Result<u32, ProviderError>;

    /// Embed `text` with the named embedding model, returning the raw vector.
    /// Backs the `similar` matcher, which embeds a reference and the
    /// actual value and compares them by cosine similarity. Providers with no
    /// embeddings endpoint (Anthropic) return
    /// [`ProviderError::FeatureUnsupported`] rather than faking a vector.
    async fn embed(&self, text: &str, model: &str) -> Result<Vec<f32>, ProviderError>;

    /// Convert a usage record into a dollar cost. The pricing table is
    /// per-impl because each provider sets its own rates.
    fn estimate_cost(&self, usage: &Usage, model: &str) -> Cost;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn provider_id_serializes_kebab_case() {
        let v = serde_json::to_value(ProviderId::Anthropic).unwrap();
        assert_eq!(v, serde_json::json!("anthropic"));
        let v = serde_json::to_value(ProviderId::Openai).unwrap();
        assert_eq!(v, serde_json::json!("openai"));
    }

    #[test]
    fn capabilities_none_disables_every_flag() {
        let c = Capabilities::none();
        assert!(!c.json_mode);
        assert!(!c.tool_use);
        assert!(!c.streaming);
        assert!(!c.structured_output);
        assert!(!c.token_count_exact);
    }

    #[test]
    fn cost_total_sums_every_line_item() {
        let c = Cost {
            input_dollars: 0.01,
            output_dollars: 0.02,
            cache_read_dollars: Some(0.003),
            cache_write_dollars: Some(0.004),
        };
        assert!((c.total() - 0.037).abs() < 1e-9);
    }

    #[test]
    fn cost_total_treats_missing_cache_fields_as_zero() {
        let c = Cost {
            input_dollars: 0.10,
            output_dollars: 0.05,
            cache_read_dollars: None,
            cache_write_dollars: None,
        };
        assert!((c.total() - 0.15).abs() < 1e-9);
    }

    #[test]
    fn retry_policy_none_disables_attempts() {
        let p = RetryPolicy::none();
        assert_eq!(p.max_retries, 0);
        assert_eq!(p.total_budget, Duration::ZERO);
    }

    #[test]
    fn retry_policy_delay_grows_exponentially() {
        let p = RetryPolicy {
            max_retries: 4,
            base_delay: Duration::from_millis(50),
            total_budget: Duration::from_secs(60),
        };
        assert_eq!(p.delay_for(0), Duration::from_millis(50));
        assert_eq!(p.delay_for(1), Duration::from_millis(100));
        assert_eq!(p.delay_for(2), Duration::from_millis(200));
        assert_eq!(p.delay_for(3), Duration::from_millis(400));
    }

    #[tokio::test]
    async fn map_reqwest_error_reports_configured_timeout_not_a_constant() {
        // A 1ms request timeout against a server that never answers in
        // time produces a real `is_timeout` error. The mapper must echo
        // the timeout it is handed (123s here), not any hardcoded value.
        let client = reqwest::Client::builder()
            .timeout(Duration::from_millis(1))
            .build()
            .expect("build client");
        // RFC 5737 TEST-NET-1 address; routable but never replies, so the
        // body read trips the request timeout rather than connecting.
        let err = client
            .get("http://192.0.2.1/")
            .send()
            .await
            .expect_err("request to a black-hole address should fail");
        let configured = Duration::from_secs(123);
        match map_reqwest_error(err, configured) {
            // Either a clean timeout or a connect failure is acceptable
            // depending on the host's routing; the timeout case is the one
            // this test pins, so retry with a connect-only address below.
            ProviderError::Timeout(d) => assert_eq!(d, configured),
            ProviderError::Transport(_) => {
                // The address connected or reset before the timeout fired;
                // re-run against a guaranteed-slow loopback delay instead.
                let server = wiremock::MockServer::start().await;
                wiremock::Mock::given(wiremock::matchers::method("GET"))
                    .respond_with(
                        wiremock::ResponseTemplate::new(200).set_delay(Duration::from_secs(30)),
                    )
                    .mount(&server)
                    .await;
                let client = reqwest::Client::builder()
                    .timeout(Duration::from_millis(20))
                    .build()
                    .expect("build client");
                let err = client
                    .get(server.uri())
                    .send()
                    .await
                    .expect_err("delayed response should trip the timeout");
                match map_reqwest_error(err, configured) {
                    ProviderError::Timeout(d) => assert_eq!(d, configured),
                    other => panic!("expected timeout, got {other:?}"),
                }
            }
            other => panic!("expected timeout or transport, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn map_reqwest_error_preserves_the_connect_branch() {
        // Connect to a port nothing is listening on. reqwest reports this
        // as an `is_connect` error; the shared mapper must keep the
        // connect tag rather than collapsing it into a bare transport
        // string (the branch OpenAI and Mistral had dropped).
        //
        // The connect_timeout is deliberately generous (5s). A closed
        // loopback port refuses in microseconds, so the refusal always wins
        // the race against the budget and the error is classified as
        // `is_connect`. A tight budget (the old 200ms) let a loaded Windows
        // runner trip the connect_timeout first, which reqwest reports as
        // `is_timeout` and the mapper turns into ProviderError::Timeout,
        // failing this test for the wrong reason.
        let client = reqwest::Client::builder()
            .connect_timeout(Duration::from_secs(5))
            .build()
            .expect("build client");
        let err = client
            // Port 1 on loopback refuses immediately on every CI host.
            .get("http://127.0.0.1:1/")
            .send()
            .await
            .expect_err("connection to a closed port should fail");
        match map_reqwest_error(err, Duration::from_secs(5)) {
            ProviderError::Transport(msg) => {
                assert!(
                    msg.starts_with("connect:"),
                    "connect failures must be tagged, got {msg:?}"
                );
            }
            other => panic!("expected a tagged transport error, got {other:?}"),
        }
    }
}
