procyon 0.1.1

//! Provider dispatch.
//!
//! An enum rather than a trait object: the two methods borrow a channel and return a future, and
//! the set of providers is closed and small, so `async_trait` boxing buys nothing here.

use std::sync::LazyLock;
use std::time::Duration;

use color_eyre::{eyre::bail, Result};
use reqwest::{Client, RequestBuilder, StatusCode};
use tokio::sync::mpsc;

use crate::agent::{ContentPart, Message, StreamOutcome, ToolDefinition};
use crate::anthropic::ClaudeClient;
use crate::config::AppConfig;
use crate::openai::OpenAiClient;

// One connection pool for the process. A `reqwest::Client` owns its pool and resolver, so building
// one per client — and a client is built per sub-agent — threw away every warm connection and
// opened a fresh TLS handshake for each persona in a party round.
static HTTP: LazyLock<Client> = LazyLock::new(Client::new);

/// The models a locally-served provider actually has, asked of the server itself.
///
/// Only for local providers. A remote catalogue changes weekly and the built-in list is a set of
/// suggestions, not an inventory — refusing a model merely because this build has not heard of it
/// would block every model released after it. A local server, by contrast, knows exactly what is
/// installed, and that answer is worth having: `/model model <anything>` used to be accepted with
/// no check at all, and the first sign of trouble was a failing turn.
///
/// `None` means the question could not be answered — the server is not running, or does not serve
/// the endpoint. That is not the same as "no models", and a caller must not read it as a refusal.
pub async fn installed_models(cfg: &AppConfig) -> Option<Vec<String>> {
    if !cfg.provider.is_local() {
        return None;
    }

    // The OpenAI-compatible listing, which Ollama and LM Studio both serve, so one path covers
    // both instead of each needing its own vendor endpoint.
    let base = cfg.resolve_base_url().ok()?;
    let url = format!("{}/models", base.trim_end_matches('/'));

    let response = HTTP
        .get(url)
        .timeout(Duration::from_secs(3))
        .send()
        .await
        .ok()?;
    if !response.status().is_success() {
        return None;
    }

    let body: serde_json::Value = response.json().await.ok()?;
    let models: Vec<String> = body
        .get("data")?
        .as_array()?
        .iter()
        .filter_map(|entry| entry.get("id")?.as_str().map(str::to_string))
        .collect();

    // An empty list is a real answer — a server running with nothing pulled — but reading it as
    // one would refuse every switch on a machine whose server has simply not been populated yet.
    if models.is_empty() {
        None
    } else {
        Some(models)
    }
}

/// Whether `model` is one the local server has, allowing for the tag Ollama leaves implicit:
/// `llama3.2` and `llama3.2:latest` are the same model, and the short form is what users type.
pub fn is_installed(models: &[String], model: &str) -> bool {
    models.iter().any(|installed| {
        installed == model
            || installed
                .strip_suffix(":latest")
                .is_some_and(|base| base == model)
    })
}

/// The context window Ollama actually loaded a model with, as opposed to the one the weights
/// advertise.
///
/// These differ by a lot and the difference is the whole problem. `qwen3:4b` declares 262144, and
/// a default `ollama serve` loads it with 4096 — `OLLAMA_CONTEXT_LENGTH` is a server setting, and
/// the OpenAI-compatible endpoint this build speaks over cannot ask for more (`options.num_ctx` is
/// not forwarded, verified against 0.20.4). So the number that matters is neither the model's nor
/// this build's: it is the server's, and `/api/ps` is where it says it.
///
/// `None` when the model is not loaded yet — nothing is allocated before the first request — or
/// when the server does not answer. Callers keep their conservative default in that case.
pub async fn ollama_context_length(cfg: &AppConfig, model: &str) -> Option<usize> {
    if cfg.provider != crate::config::Provider::Ollama {
        return None;
    }

    // `/api/ps` sits beside `/v1`, not under it.
    let base = cfg.resolve_base_url().ok()?;
    let root = base.trim_end_matches('/').strip_suffix("/v1")?;

    let response = HTTP
        .get(format!("{}/api/ps", root))
        .timeout(Duration::from_secs(2))
        .send()
        .await
        .ok()?;
    if !response.status().is_success() {
        return None;
    }

    let body: serde_json::Value = response.json().await.ok()?;
    body.get("models")?
        .as_array()?
        .iter()
        .find(|entry| {
            entry
                .get("name")
                .and_then(|n| n.as_str())
                .is_some_and(|name| is_installed(&[name.to_string()], model))
        })?
        .get("context_length")?
        .as_u64()
        .map(|n| n as usize)
}

/// Returns a handle on the shared HTTP client. Cloning shares the underlying pool.
pub fn http_client() -> Client {
    HTTP.clone()
}

/// Attempts one request makes before its failure is reported.
const MAX_ATTEMPTS: u32 = 3;

/// Ceiling on a server-supplied `Retry-After`, so a provider asking for a ten minute wait does not
/// hold the whole session.
const MAX_RETRY_AFTER: Duration = Duration::from_secs(20);

/// Whether a status is worth trying again.
///
/// Rate limits and 5xx are the provider saying "not now"; every 4xx below 429 is the request
/// itself being wrong, and repeating it only wastes the quota. 408 is a gateway giving up on a
/// request the model may never have seen.
fn is_transient(status: StatusCode) -> bool {
    status == StatusCode::TOO_MANY_REQUESTS
        || status == StatusCode::REQUEST_TIMEOUT
        || status.is_server_error()
}

/// Whether a transport failure is worth trying again. A refused connection or a dropped socket may
/// be the network; a body that could not be built is us.
fn is_transient_transport(error: &reqwest::Error) -> bool {
    error.is_timeout() || error.is_connect() || error.is_request()
}

fn backoff(attempt: u32) -> Duration {
    // 500ms, 1s. Short: a turn is interactive, and the point is to survive a blip rather than to
    // wait out a sustained outage.
    Duration::from_millis(500 * 2u64.pow(attempt.saturating_sub(1)))
}

/// Sends a request, retrying the failures that are the provider's rather than ours.
///
/// Every non-2xx used to end the turn outright, which meant a single `overloaded_error` threw away
/// the whole turn — including every tool result already produced and paid for. `build` is a
/// factory rather than a `RequestBuilder` because a retry needs a fresh body.
pub async fn send_with_retry<F>(build: F) -> Result<reqwest::Response>
where
    F: Fn() -> RequestBuilder,
{
    let mut attempt = 0;

    loop {
        attempt += 1;
        let last = attempt >= MAX_ATTEMPTS;

        match build().send().await {
            Ok(response) if response.status().is_success() => return Ok(response),
            Ok(response) => {
                let status = response.status();
                if !is_transient(status) || last {
                    // Read after the retry decision: consuming the body moves the response, and a
                    // retryable status does not need its text.
                    let body = response.text().await.unwrap_or_default();
                    bail!("API error {}: {}", status, body);
                }
                let wait = retry_after(&response).unwrap_or_else(|| backoff(attempt));
                crate::diag::warn(format!(
                    "provider returned {}, retrying in {:?} (attempt {} of {})",
                    status, wait, attempt, MAX_ATTEMPTS
                ));
                tokio::time::sleep(wait).await;
            }
            Err(e) => {
                if !is_transient_transport(&e) || last {
                    return Err(e.into());
                }
                let wait = backoff(attempt);
                crate::diag::warn(format!(
                    "request failed ({}), retrying in {:?} (attempt {} of {})",
                    e, wait, attempt, MAX_ATTEMPTS
                ));
                tokio::time::sleep(wait).await;
            }
        }
    }
}

/// Honours the provider's own pacing when it states one. Only the delta-seconds form is read; the
/// HTTP-date form is rare here and a wrong parse would be worse than a backoff.
fn retry_after(response: &reqwest::Response) -> Option<Duration> {
    let seconds: u64 = response
        .headers()
        .get(reqwest::header::RETRY_AFTER)?
        .to_str()
        .ok()?
        .trim()
        .parse()
        .ok()?;

    Some(Duration::from_secs(seconds).min(MAX_RETRY_AFTER))
}

pub enum LlmClient {
    Anthropic(ClaudeClient),
    OpenAiCompatible(OpenAiClient),
}

impl LlmClient {
    /// Builds the client the config selects, resolving the credential for that provider.
    pub fn from_config(config: &AppConfig) -> Result<Self> {
        let api_key = config.get_api_key()?;

        // Dispatch is on the protocol: every provider but Anthropic speaks the OpenAI dialect,
        // differing only in the endpoint the config already resolved.
        Ok(if config.provider.is_anthropic() {
            Self::Anthropic(
                ClaudeClient::new(api_key)
                    .with_model(config.default_model.clone())
                    .with_max_tokens(config.max_tokens),
            )
        } else {
            Self::OpenAiCompatible(OpenAiClient::new(
                api_key,
                config.resolve_base_url()?,
                config.default_model.clone(),
                config.max_tokens,
            ))
        })
    }

    pub async fn send_message(
        &self,
        history: &[Message],
        tools: Option<&[ToolDefinition]>,
        system: Option<&str>,
    ) -> Result<Vec<ContentPart>> {
        match self {
            Self::Anthropic(client) => client.send_message(history, tools, system).await,
            Self::OpenAiCompatible(client) => client.send_message(history, tools, system).await,
        }
    }

    pub async fn send_message_streaming(
        &self,
        history: &[Message],
        tools: Option<&[ToolDefinition]>,
        system: Option<&str>,
        update_tx: &mpsc::UnboundedSender<String>,
    ) -> Result<StreamOutcome> {
        match self {
            Self::Anthropic(client) => {
                client
                    .send_message_streaming(history, tools, system, update_tx)
                    .await
            }
            Self::OpenAiCompatible(client) => {
                client
                    .send_message_streaming(history, tools, system, update_tx)
                    .await
            }
        }
    }
}

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

    #[test]
    fn an_exact_name_is_installed() {
        let installed = vec!["qwen3:4b-instruct-2507-q4_K_M".to_string()];
        assert!(is_installed(&installed, "qwen3:4b-instruct-2507-q4_K_M"));
    }

    // Ollama lists `llama3.2:latest`, and nobody types the tag.
    #[test]
    fn the_implicit_latest_tag_still_matches() {
        let installed = vec!["llama3.2:latest".to_string()];
        assert!(is_installed(&installed, "llama3.2"));
        assert!(is_installed(&installed, "llama3.2:latest"));
    }

    // The other direction must not hold: `llama3.2` installed does not mean `llama3.2:8b` is.
    #[test]
    fn a_different_tag_is_a_different_model() {
        let installed = vec!["llama3.1:8b".to_string(), "llama3.2".to_string()];
        assert!(!is_installed(&installed, "llama3.1"));
        assert!(!is_installed(&installed, "llama3.2:70b"));
        assert!(!is_installed(&installed, "mistral"));
    }

    #[tokio::test]
    async fn a_remote_provider_is_never_second_guessed() {
        // A vendor catalogue changes weekly and the built-in list is suggestions, not an
        // inventory. Refusing an unlisted model would block every model released after this build.
        let cfg = AppConfig {
            provider: Provider::Anthropic,
            ..AppConfig::default()
        };
        assert!(installed_models(&cfg).await.is_none());
    }

    #[tokio::test]
    async fn a_server_that_is_not_running_answers_nothing_rather_than_nothing_installed() {
        // Port 1 has no listener. `None` has to mean "could not ask", so that a switch is allowed
        // through: refusing every model because the server is down would be worse than the bug.
        let cfg = AppConfig {
            provider: Provider::Ollama,
            base_url: Some("http://127.0.0.1:1/v1".to_string()),
            ..AppConfig::default()
        };
        assert!(installed_models(&cfg).await.is_none());
    }
}

#[cfg(test)]
mod ollama_probe {
    use super::*;
    use crate::config::Provider;

    /// The number that matters is the server's, not the model's. `qwen3:4b` advertises 262144 and
    /// a default `ollama serve` loads it with 4096, so this must read what was actually allocated.
    ///
    /// Needs a second server started with a window that differs from the default, which is the
    /// only way to tell detection from the hardcoded fallback:
    ///
    ///   OLLAMA_HOST=127.0.0.1:11435 OLLAMA_CONTEXT_LENGTH=32768 \
    ///     OLLAMA_MODELS=/usr/share/ollama/.ollama/models ollama serve
    ///
    /// then load a model on it once. cargo test reads_the_window_the_server_loaded -- --ignored
    #[tokio::test]
    #[ignore]
    async fn reads_the_window_the_server_loaded() {
        let cfg = AppConfig {
            provider: Provider::Ollama,
            base_url: Some("http://127.0.0.1:11435/v1".to_string()),
            ..AppConfig::default()
        };

        let found = ollama_context_length(&cfg, "qwen3:4b-instruct-2507-q4_K_M").await;
        assert_eq!(
            found,
            Some(32_768),
            "detection must follow the server, not the 4096 default"
        );
    }
}