captchaforge 0.2.40

//! Model download and cache management for vision inference.
//!
//! Models are stored under the platform cache directory
//! (`dirs::cache_dir()/captchaforge/models/`: `~/.cache/captchaforge/models/`
//! on Linux, honoring `$XDG_CACHE_HOME`) and downloaded from Hugging Face or
//! GitHub releases on first use. The CRNN text model has no download URL; it is
//! produced locally by `scripts/train_crnn.py`, whose default `--output-dir`
//! resolves to this same cache path.

use anyhow::{Context, Result};
use std::path::{Path, PathBuf};
use tracing::info;

/// Canonical model cache directory.
fn cache_dir() -> PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(std::env::temp_dir)
        .join("captchaforge")
        .join("models")
}

/// Unique identifier for each downloadable model.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ModelId {
    /// YOLOv8 nano: COCO pre-trained object detection.
    /// ~6 MB, fast inference, 80 COCO classes.
    YoloV8n,
    /// CRNN text recognition (placeholder (needs trained model)).
    CrnnText,
}

impl ModelId {
    fn repo_path(self) -> (&'static str, &'static str) {
        match self {
            Self::YoloV8n => (
                "https://github.com/ultralytics/assets/releases/download/v8.4.0/yolov8n.onnx",
                "yolov8n.onnx",
            ),
            Self::CrnnText => ("", "crnn_text.onnx"),
        }
    }

    fn cache_subdir(self) -> &'static str {
        match self {
            Self::YoloV8n => "yolov8n",
            Self::CrnnText => "crnn_text",
        }
    }
}

/// Manages model download and local cache.
#[derive(Debug, Clone)]
pub struct ModelHub {
    root: PathBuf,
}

impl Default for ModelHub {
    fn default() -> Self {
        Self::new()
    }
}

impl ModelHub {
    pub fn new() -> Self {
        Self { root: cache_dir() }
    }

    /// Override the cache root (useful for tests).
    pub fn with_root(root: impl Into<PathBuf>) -> Self {
        Self { root: root.into() }
    }

    /// Return the local path for a model, downloading it if necessary.
    pub async fn resolve(&self, id: ModelId) -> Result<PathBuf> {
        let subdir = self.root.join(id.cache_subdir());
        let (_, filename) = id.repo_path();
        let path = subdir.join(filename);

        if path.exists() {
            info!(model = ?id, path = %path.display(), "model cached");
            return Ok(path);
        }

        let (url, _) = id.repo_path();
        if url.is_empty() {
            anyhow::bail!(
                "model {:?} has no download URL, must be provided manually at {}",
                id,
                path.display()
            );
        }

        info!(model = ?id, url, "downloading model");
        self.download(url, &path).await?;
        Ok(path)
    }

    /// Return the local path for a model, *without* downloading.
    /// Returns `None` if the model is not cached.
    pub fn local_path(&self, id: ModelId) -> Option<PathBuf> {
        let (_, filename) = id.repo_path();
        let path = self.root.join(id.cache_subdir()).join(filename);
        path.exists().then_some(path)
    }

    async fn download(&self, url: &str, dest: &Path) -> Result<()> {
        std::fs::create_dir_all(dest.parent().unwrap())?;

        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(300))
            .build()
            .context("build download client")?;

        let resp = client
            .get(url)
            .send()
            .await
            .with_context(|| format!("download request failed: {url}"))?;

        if !resp.status().is_success() {
            anyhow::bail!("download returned {}: {}", resp.status(), url);
        }

        let bytes = resp.bytes().await.context("read download body")?;

        tokio::fs::write(dest, bytes)
            .await
            .with_context(|| format!("write model to {}", dest.display()))?;

        info!(path = %dest.display(), "model downloaded");
        Ok(())
    }
}

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

    #[test]
    fn cache_dir_is_reasonable() {
        let d = cache_dir();
        assert!(d.to_string_lossy().contains("captchaforge"));
    }

    #[test]
    fn hub_local_path_missing() {
        let hub = ModelHub::with_root("/nonexistent/captchaforge/models");
        assert!(hub.local_path(ModelId::YoloV8n).is_none());
    }
}