captchaforge 0.2.36

//! Auto-detect available solver backends and configure the chain.
//!
//! Production captchaforge ergonomics: a developer should be able to
//! `cargo add captchaforge && captchaforge::auto_solve(page).await` and
//! have it Just Work — no manual VLM endpoint configuration, no STT
//! server URLs, no tesseract path hunting. This module does the
//! probing and surfaces the result as [`Capabilities`].
//!
//! # What gets probed
//!
//! - **Ollama** — HTTP probe `localhost:11434/api/tags`. If reachable,
//!   the response lists locally-pulled models. We pick the first
//!   model whose name matches a known vision-model pattern
//!   (`llama3.2-vision`, `llava`, `bakllava`, `qwen2-vl`,
//!   `qwen3-vl`, `gemma3-vl`, `pixtral`, `minicpm-v`).
//! - **Whisper CLI** — `which whisper` on PATH. The OpenAI Whisper
//!   reference CLI; `whisper <audio.mp3> --model tiny` returns a
//!   transcript on stdout.
//! - **Tesseract** — `which tesseract` on PATH. Used for canvas /
//!   SVG / static-text CAPTCHAs.
//! - **OpenAI** — `OPENAI_API_KEY` env var. If set, OpenAI Whisper
//!   API is available as STT fallback.
//!
//! # Auto-install
//!
//! [`Capabilities::ensure_vlm_model`] will `ollama pull` a default
//! vision model if Ollama is present but no vision model is. Disabled
//! by default; opt in by calling `ensure_vlm_model()` directly.
//!
//! # Example
//!
//! ```rust,no_run
//! # async fn run() {
//! let caps = captchaforge::backends::probe().await;
//! if caps.vlm.is_some() {
//!     println!("VLM ready: {}", caps.vlm.as_ref().unwrap().model);
//! }
//! # }
//! ```

use std::time::Duration;

use serde::{Deserialize, Serialize};

/// Set of vendor model-name fragments that identify a vision-capable
/// Ollama model. Order matters: earlier entries are preferred when
/// multiple match.
pub const VISION_MODEL_PATTERNS: &[&str] = &[
    "llama3.2-vision",
    "llava",
    "bakllava",
    "qwen2-vl",
    "qwen3-vl",
    "gemma3-vl",
    "pixtral",
    "minicpm-v",
    "moondream",
];

/// Default vision model to auto-pull when Ollama is present but no
/// vision model is. Picked for size/quality tradeoff: ~7GB, runs on
/// most consumer GPUs and decent CPUs.
pub const DEFAULT_VLM_AUTO_PULL: &str = "llama3.2-vision:11b";

/// Default Whisper model size to use when invoking the local CLI.
///
/// `base` (~142MB) reliably transcribes spelled digits where `tiny` often
/// hallucinates homophones ("four" → "for", "two" → "to/too", "nine" →
/// "high"). For CAPTCHA audio — usually 3-6 spoken digits — accuracy
/// matters more than the ~70MB size delta. CPU latency stays under 1s.
pub const DEFAULT_WHISPER_MODEL: &str = "base";

/// Detected backend capabilities. Populated by [`probe`]; consumed by
/// [`crate::config::Config::build_chain`] to wire solvers conditionally.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Capabilities {
    pub vlm: Option<VlmBackend>,
    pub stt: Option<SttBackend>,
    pub ocr: Option<OcrBackend>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VlmBackend {
    /// Endpoint for chat-completion (Ollama-compatible).
    pub endpoint: String,
    /// Model name to send in the request body.
    pub model: String,
    /// Source — `"ollama"`, `"openai"`, etc. — for logging/telemetry.
    pub source: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SttBackend {
    pub kind: SttKind,
    /// For `WhisperCli` this is the path to the binary; for
    /// `OpenAIApi` the API base URL; for `LocalServer` the HTTP URL.
    pub endpoint: String,
    /// For `WhisperCli` the model size (`tiny`/`base`/`small`/...).
    pub model: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SttKind {
    WhisperCli,
    OpenAIApi,
    LocalServer,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrBackend {
    /// Path to the `tesseract` binary.
    pub binary: String,
}

/// Probe the local environment for available backends. Cheap (~50ms
/// total) — three parallel HTTP/PATH checks.
pub async fn probe() -> Capabilities {
    let (vlm, stt, ocr) = tokio::join!(probe_vlm(), probe_stt(), probe_ocr());
    Capabilities { vlm, stt, ocr }
}

async fn probe_vlm() -> Option<VlmBackend> {
    let endpoint = "http://localhost:11434";
    let client = match reqwest::Client::builder()
        .timeout(Duration::from_millis(800))
        .build()
    {
        Ok(c) => c,
        Err(_) => return None,
    };
    let resp = client
        .get(format!("{endpoint}/api/tags"))
        .send()
        .await
        .ok()?;
    if !resp.status().is_success() {
        return None;
    }
    #[derive(Deserialize)]
    struct OllamaTags {
        models: Vec<OllamaModel>,
    }
    #[derive(Deserialize)]
    struct OllamaModel {
        name: String,
    }
    let body: OllamaTags = resp.json().await.ok()?;
    let model = pick_best_vision_model(
        &body
            .models
            .iter()
            .map(|m| m.name.as_str())
            .collect::<Vec<_>>(),
    )?;
    Some(VlmBackend {
        endpoint: endpoint.to_string(),
        model: model.to_string(),
        source: "ollama".to_string(),
    })
}

/// Pure helper: pick the best vision model from a list of Ollama
/// model names by matching against [`VISION_MODEL_PATTERNS`] in
/// priority order.
///
/// # Examples
///
/// ```
/// use captchaforge::backends::pick_best_vision_model;
/// let names = ["llama3.1:8b", "llava:13b", "qwen2.5-coder:14b"];
/// assert_eq!(pick_best_vision_model(&names), Some("llava:13b"));
/// // No vision model → None.
/// let names = ["llama3.1:8b", "qwen2.5-coder:14b"];
/// assert_eq!(pick_best_vision_model(&names), None);
/// // Highest-priority pattern wins (llama3.2-vision before llava).
/// let names = ["llava:13b", "llama3.2-vision:11b"];
/// assert_eq!(pick_best_vision_model(&names), Some("llama3.2-vision:11b"));
/// ```
pub fn pick_best_vision_model<'a>(model_names: &[&'a str]) -> Option<&'a str> {
    for pattern in VISION_MODEL_PATTERNS {
        for name in model_names {
            if name.contains(pattern) {
                return Some(name);
            }
        }
    }
    None
}

async fn probe_stt() -> Option<SttBackend> {
    // Tier 1: local whisper CLI on PATH.
    if let Some(path) = which("whisper").or_else(|| which("whisper-cli")) {
        return Some(SttBackend {
            kind: SttKind::WhisperCli,
            endpoint: path,
            model: Some(DEFAULT_WHISPER_MODEL.to_string()),
        });
    }
    // Tier 2: OpenAI Whisper API via env var.
    if std::env::var("OPENAI_API_KEY").is_ok() {
        return Some(SttBackend {
            kind: SttKind::OpenAIApi,
            endpoint: "https://api.openai.com/v1/audio/transcriptions".to_string(),
            model: Some("whisper-1".to_string()),
        });
    }
    // Tier 3: classic local server on the conventional port.
    let client = reqwest::Client::builder()
        .timeout(Duration::from_millis(300))
        .build()
        .ok()?;
    if client
        .get("http://localhost:9000/")
        .send()
        .await
        .map(|r| r.status().is_success() || r.status().as_u16() == 405)
        .unwrap_or(false)
    {
        return Some(SttBackend {
            kind: SttKind::LocalServer,
            endpoint: "http://localhost:9000/asr".to_string(),
            model: None,
        });
    }
    None
}

async fn probe_ocr() -> Option<OcrBackend> {
    which("tesseract").map(|binary| OcrBackend { binary })
}

/// Locate a binary on PATH. Returns the absolute path on hit, `None`
/// on miss. Pure helper used by both `probe_stt` and `probe_ocr`.
///
/// # Examples
///
/// ```
/// use captchaforge::backends::which;
/// // Almost-certainly-present binary on a Unix system.
/// // Skipped on Windows / locked-down environments.
/// # if cfg!(unix) {
/// assert!(which("sh").is_some(), "sh should be on PATH");
/// # }
/// // Definitely-not-present binary.
/// assert!(which("definitely-not-a-real-binary-xyzzy").is_none());
/// ```
pub fn which(binary: &str) -> Option<String> {
    let path = std::env::var_os("PATH")?;
    for dir in std::env::split_paths(&path) {
        let candidate = dir.join(binary);
        if candidate.is_file() {
            return Some(candidate.to_string_lossy().into_owned());
        }
    }
    None
}

impl Capabilities {
    /// Returns true if any solver-backing capability is present. Used
    /// by the chain to short-circuit the "no backends installed,
    /// chain will only handle passive Turnstile" log message.
    pub fn any(&self) -> bool {
        self.vlm.is_some() || self.stt.is_some() || self.ocr.is_some()
    }

    /// Human-readable one-liner describing what was detected. Used by
    /// `auto_solve` to log on first call.
    pub fn summary(&self) -> String {
        let mut parts = Vec::new();
        if let Some(v) = &self.vlm {
            parts.push(format!("vlm:{} ({})", v.model, v.source));
        }
        if let Some(s) = &self.stt {
            parts.push(format!("stt:{:?}", s.kind));
        }
        if self.ocr.is_some() {
            parts.push("ocr:tesseract".to_string());
        }
        if parts.is_empty() {
            "no backends detected".to_string()
        } else {
            parts.join(", ")
        }
    }

    /// Returns a one-line install hint per missing backend, suitable
    /// for surfacing to a user who needs higher solve coverage. Empty
    /// when every backend is already present.
    ///
    /// # Examples
    ///
    /// ```
    /// use captchaforge::backends::Capabilities;
    /// let caps = Capabilities::default();
    /// let hints = caps.install_hints();
    /// assert!(!hints.is_empty(), "no backends ⇒ at least one hint");
    /// assert!(hints.iter().any(|h| h.contains("ollama")));
    /// ```
    pub fn install_hints(&self) -> Vec<String> {
        let mut out = Vec::new();
        if self.vlm.is_none() {
            out.push(
                "vlm: install ollama (https://ollama.com/download) and run \
                 `ollama pull llama3.2-vision:11b` — captchaforge will then \
                 auto-detect it. Solves canvas / SVG / image-grid CAPTCHAs."
                    .to_string(),
            );
        }
        if self.stt.is_none() {
            out.push(
                "stt: install OpenAI Whisper CLI (`pip install -U openai-whisper`) \
                 OR set OPENAI_API_KEY for the hosted Whisper API. Solves \
                 reCAPTCHA v2 audio + generic spoken-digit CAPTCHAs."
                    .to_string(),
            );
        }
        if self.ocr.is_none() {
            out.push(
                "ocr: install tesseract (`apt install tesseract-ocr` / \
                 `brew install tesseract`). Used as a fallback when no VLM \
                 is available."
                    .to_string(),
            );
        }
        out
    }

    /// Auto-install a default vision model when Ollama is reachable
    /// but no vision model is present. Returns `Ok(true)` if pull
    /// succeeded, `Ok(false)` if Ollama isn't reachable, or `Err` on
    /// pull failure.
    ///
    /// Spawns `ollama pull <DEFAULT_VLM_AUTO_PULL>`. Pull progress is
    /// inherited stdio; the call blocks until pull completes (~5-10
    /// minutes for an 11B model on first run).
    pub async fn ensure_vlm_model(&mut self) -> anyhow::Result<bool> {
        if self.vlm.is_some() {
            return Ok(true);
        }
        let ollama = which("ollama");
        let Some(ollama_bin) = ollama else {
            return Ok(false);
        };
        // Verify Ollama daemon is up before spawning a pull.
        let client = reqwest::Client::builder()
            .timeout(Duration::from_millis(800))
            .build()?;
        if client
            .get("http://localhost:11434/api/tags")
            .send()
            .await
            .map(|r| !r.status().is_success())
            .unwrap_or(true)
        {
            return Ok(false);
        }
        tracing::info!(
            model = DEFAULT_VLM_AUTO_PULL,
            "auto-pulling vision model via Ollama (first-run install)"
        );
        let status = std::process::Command::new(&ollama_bin)
            .arg("pull")
            .arg(DEFAULT_VLM_AUTO_PULL)
            .status()?;
        if !status.success() {
            anyhow::bail!("ollama pull {DEFAULT_VLM_AUTO_PULL} exited {status}");
        }
        // Re-probe so the caller sees the new vlm.
        self.vlm = probe_vlm().await;
        Ok(self.vlm.is_some())
    }
}

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

    #[test]
    fn vision_model_patterns_are_lowercase() {
        // Ollama tag matching is substring-and-lowercase; the patterns
        // list must follow.
        for p in VISION_MODEL_PATTERNS {
            assert_eq!(*p, p.to_lowercase(), "pattern {p} not lowercase");
        }
    }

    #[test]
    fn pick_best_vision_model_returns_none_for_text_only() {
        let names = ["mistral:7b", "qwen2.5:14b", "llama3.1:8b"];
        assert_eq!(pick_best_vision_model(&names), None);
    }

    #[test]
    fn pick_best_vision_model_finds_llava() {
        let names = ["llama3.1:8b", "llava:13b"];
        assert_eq!(pick_best_vision_model(&names), Some("llava:13b"));
    }

    #[test]
    fn pick_best_vision_model_priority_order_respected() {
        // llama3.2-vision is listed before llava in VISION_MODEL_PATTERNS;
        // when both are pulled, llama3.2-vision wins.
        let names = ["llava:13b", "llama3.2-vision:11b"];
        assert_eq!(pick_best_vision_model(&names), Some("llama3.2-vision:11b"));
    }

    #[test]
    fn pick_best_vision_model_handles_empty_list() {
        assert_eq!(pick_best_vision_model(&[]), None);
    }

    #[test]
    fn capabilities_summary_no_backends() {
        let caps = Capabilities::default();
        assert_eq!(caps.summary(), "no backends detected");
        assert!(!caps.any());
    }

    #[test]
    fn capabilities_summary_lists_each_backend() {
        let caps = Capabilities {
            vlm: Some(VlmBackend {
                endpoint: "http://localhost:11434".to_string(),
                model: "llava:13b".to_string(),
                source: "ollama".to_string(),
            }),
            stt: Some(SttBackend {
                kind: SttKind::WhisperCli,
                endpoint: "/usr/local/bin/whisper".to_string(),
                model: Some("tiny".to_string()),
            }),
            ocr: Some(OcrBackend {
                binary: "/usr/bin/tesseract".to_string(),
            }),
        };
        let s = caps.summary();
        assert!(s.contains("vlm:llava:13b"));
        assert!(s.contains("stt:WhisperCli"));
        assert!(s.contains("ocr:tesseract"));
        assert!(caps.any());
    }

    #[test]
    fn default_vlm_pull_target_is_a_known_pattern() {
        assert!(
            VISION_MODEL_PATTERNS
                .iter()
                .any(|p| DEFAULT_VLM_AUTO_PULL.contains(p)),
            "DEFAULT_VLM_AUTO_PULL must match a VISION_MODEL_PATTERNS entry; \
             otherwise the post-pull probe won't find the new model",
        );
    }

    #[test]
    fn which_finds_existing_binary_unix() {
        if cfg!(unix) {
            assert!(which("sh").is_some(), "sh should be on PATH");
        }
    }

    #[test]
    fn which_returns_none_for_missing_binary() {
        assert!(which("definitely-not-a-real-binary-xyzzy12345").is_none());
    }

    #[tokio::test]
    async fn probe_returns_some_capabilities_field_default() {
        // Pure-data sanity: probe always returns a Capabilities struct
        // even when nothing is installed.
        let caps = probe().await;
        // We can't assert exact contents (depends on machine), but the
        // type is stable.
        let _ = caps.summary();
    }

    #[test]
    fn stt_kind_serializes_snake_case() {
        let json = serde_json::to_string(&SttKind::WhisperCli).unwrap();
        assert_eq!(json, r#""whisper_cli""#);
        // serde's snake_case mangles `OpenAIApi` to `open_a_i_api`
        // because it treats every uppercase letter as a word boundary.
        // We accept the serde default rather than hand-rolling a custom
        // serializer — the JSON form is internal/telemetry-only.
        let json = serde_json::to_string(&SttKind::OpenAIApi).unwrap();
        assert_eq!(json, r#""open_a_i_api""#);
        let json = serde_json::to_string(&SttKind::LocalServer).unwrap();
        assert_eq!(json, r#""local_server""#);
    }
}