//! Ollama provider impl.
//!
//! Targets the local Ollama HTTP server (default
//! `http://localhost:11434/api/chat`). No auth header, no rate-limit
//! semantics in the wire protocol, so the retry policy is disabled by
//! default. Cost is reported as zero because the user pays for the GPU
//! once, not per token.

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, Usage,
};
use crate::eval::providers::parse_embedding;
use crate::eval::providers::retry::run_with_retry;

const DEFAULT_BASE_URL: &str = "http://localhost:11434";

/// Ollama provider client.
pub struct OllamaProvider {
    client: reqwest::Client,
    base_url: String,
    retry: RetryPolicy,
    request_timeout: Duration,
}

/// Config bag for [`OllamaProvider::new`].
#[derive(Debug, Clone)]
pub struct OllamaConfig {
    /// Override base URL (test-server port, remote Ollama gateway).
    pub base_url: Option<String>,
    /// Body-read timeout. Default 60s; slow local models routinely run
    /// past 30s.
    pub request_timeout: Duration,
    /// Connect timeout. Default 5s; local connection refused should be
    /// the fast fail signal, not a 30s hang.
    pub connect_timeout: Duration,
    /// HTTP / HTTPS proxy policy.
    pub proxy: crate::network::ProxyConfig,
}

impl Default for OllamaConfig {
    fn default() -> Self {
        Self {
            base_url: None,
            request_timeout: Duration::from_secs(60),
            connect_timeout: Duration::from_secs(5),
            proxy: crate::network::ProxyConfig::default(),
        }
    }
}

impl OllamaProvider {
    /// Build a new client.
    pub fn new(config: OllamaConfig) -> 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,
            base_url: config.base_url.unwrap_or_else(|| DEFAULT_BASE_URL.into()),
            retry: RetryPolicy::none(),
            request_timeout: config.request_timeout,
        })
    }
}

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

    fn supported_capabilities(&self) -> Capabilities {
        Capabilities {
            json_mode: true,
            tool_use: false,
            streaming: true,
            structured_output: false,
            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!("{}/api/chat", self.base_url);
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .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> {
        // Ollama exposes no token-count endpoint; estimate with cl100k.
        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!("{}/api/embeddings", self.base_url);
        let body = json!({ "model": model, "prompt": text });
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_embedding(&response_json, "/embedding")
    }

    fn estimate_cost(&self, _usage: &Usage, _model: &str) -> Cost {
        Cost::default()
    }
}

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,
        "stream": false,
    });
    if options.json_mode {
        body["format"] = json!("json");
    }
    if let Some(t) = options.temperature {
        body["options"] = json!({"temperature": t});
    }
    body
}

async fn handle_response(resp: reqwest::Response) -> Result<Value, ProviderError> {
    let status = resp.status();
    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() {
        404 => Err(ProviderError::InvalidRequest(format!(
            "model not found: {body}"
        ))),
        400 | 422 => Err(ProviderError::InvalidRequest(body)),
        500..=599 => Err(ProviderError::ProviderInternal(body)),
        _ => Err(ProviderError::ProviderInternal(format!(
            "unexpected status {status}: {body}"
        ))),
    }
}

#[derive(Deserialize)]
struct OllamaUsage {
    #[serde(default)]
    prompt_eval_count: u32,
    #[serde(default)]
    eval_count: 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 text = value
        .pointer("/message/content")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_owned();
    let usage_obj: OllamaUsage = serde_json::from_value(value.clone())
        .map_err(|e| ProviderError::MalformedResponse(format!("usage parse: {e}")))?;
    let finish = value
        .get("done_reason")
        .and_then(|v| v.as_str())
        .unwrap_or("stop");
    Ok(Response {
        content: ResponseContent::Text { text },
        finish_reason: match finish {
            "stop" => FinishReason::Stop,
            "length" => FinishReason::MaxTokens,
            _ => FinishReason::Other,
        },
        usage: Usage {
            input_tokens: usage_obj.prompt_eval_count,
            output_tokens: usage_obj.eval_count,
            cache_read_tokens: None,
            cache_write_tokens: None,
        },
        model,
    })
}

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

    fn provider_for(server: &MockServer) -> OllamaProvider {
        OllamaProvider::new(OllamaConfig {
            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_call_returns_text() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/api/chat"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "model": "llama3",
                "message": {"role": "assistant", "content": "local response"},
                "done": true,
                "done_reason": "stop",
                "prompt_eval_count": 7,
                "eval_count": 3,
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let messages = vec![Message {
            role: Role::User,
            content: "hi".into(),
        }];
        let options = CompletionOptions {
            model: "llama3".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&messages, &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "local response"),
            other => panic!("expected text, got {other:?}"),
        }
        assert_eq!(response.usage.input_tokens, 7);
        assert_eq!(response.usage.output_tokens, 3);
    }

    #[tokio::test]
    async fn model_not_found_maps_to_invalid_request() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/api/chat"))
            .respond_with(
                ResponseTemplate::new(404)
                    .set_body_string("model 'nope' not found, try pulling it first"),
            )
            .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(_)));
    }

    #[tokio::test]
    async fn cost_is_zero_for_local_runs() {
        let provider = OllamaProvider::new(OllamaConfig {
            base_url: Some("http://localhost".into()),
            ..Default::default()
        })
        .unwrap();
        let cost = provider.estimate_cost(
            &Usage {
                input_tokens: 1_000_000,
                output_tokens: 1_000_000,
                ..Default::default()
            },
            "llama3",
        );
        assert!((cost.total() - 0.0).abs() < f64::EPSILON);
    }
}
