visual-cortex-ocr-onnx 0.8.0

use std::path::{Path, PathBuf};
use std::time::Duration;

use sha2::{Digest, Sha256};

use crate::error::OcrError;

/// Generous ceiling for a ~10 MB model download; surfaces hung connections
/// as Download errors instead of blocking forever.
const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(300);

/// One pinned, checksummed model artifact.
#[derive(Debug, Clone, Copy)]
pub struct ModelFile {
    pub filename: &'static str,
    pub url: &'static str,
    pub sha256: &'static str,
}

/// PP-OCRv4 mobile detection model (multilingual; finds text boxes).
pub const DET_MODEL: ModelFile = ModelFile {
    filename: "ch_PP-OCRv4_det_infer.onnx",
    url: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv4/ch_PP-OCRv4_det_infer.onnx",
    sha256: "d2a7720d45a54257208b1e13e36a8479894cb74155a5efe29462512d42f49da9",
};

/// PP-OCRv3 English mobile recognition model (reads each box).
pub const REC_MODEL: ModelFile = ModelFile {
    filename: "en_PP-OCRv3_rec_infer.onnx",
    url: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv3/en_PP-OCRv3_rec_infer.onnx",
    sha256: "ef7abd8bd3629ae57ea2c28b425c1bd258a871b93fd2fe7c433946ade9b5d9ea",
};

/// English recognition charset (95 entries; CTC classes = blank + 95 + space).
pub const REC_DICT: ModelFile = ModelFile {
    filename: "en_dict.txt",
    url: "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/v2.7.0/ppocr/utils/en_dict.txt",
    sha256: "5662df9d2d03f0e8ca0d3b0649d6acbab904b6a14b3d3521463c71c37c668ce3",
};

pub(crate) fn sha256_hex(bytes: &[u8]) -> String {
    let digest = Sha256::digest(bytes);
    digest.iter().map(|b| format!("{b:02x}")).collect()
}

/// The default cache directory: `<platform cache dir>/visual-cortex`.
pub(crate) fn default_cache_dir() -> PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(std::env::temp_dir)
        .join("visual-cortex")
}

/// Return the cached path for `file`, downloading and verifying if needed.
/// A cached file failing its checksum is deleted and re-downloaded once.
pub(crate) async fn ensure_cached(file: &ModelFile, dir: &Path) -> Result<PathBuf, OcrError> {
    let path = dir.join(file.filename);
    if path.exists() {
        let bytes = tokio::fs::read(&path)
            .await
            .map_err(|e| OcrError::Download {
                url: file.url.to_string(),
                reason: format!("cache read failed: {e}"),
            })?;
        if sha256_hex(&bytes) == file.sha256 {
            return Ok(path);
        }
        tracing::warn!(
            "checksum mismatch for cached {}; re-downloading",
            path.display()
        );
        let _ = tokio::fs::remove_file(&path).await;
    }

    tokio::fs::create_dir_all(dir)
        .await
        .map_err(|e| OcrError::Download {
            url: file.url.to_string(),
            reason: format!("cache dir creation failed: {e}"),
        })?;
    let client = reqwest::Client::builder()
        .timeout(DOWNLOAD_TIMEOUT)
        .build()
        .map_err(|e| OcrError::Download {
            url: file.url.to_string(),
            reason: format!("http client build failed: {e}"),
        })?;
    let response = client
        .get(file.url)
        .send()
        .await
        .map_err(|e| OcrError::Download {
            url: file.url.to_string(),
            reason: e.to_string(),
        })?;
    if !response.status().is_success() {
        return Err(OcrError::Download {
            url: file.url.to_string(),
            reason: format!("HTTP {}", response.status()),
        });
    }
    let bytes = response.bytes().await.map_err(|e| OcrError::Download {
        url: file.url.to_string(),
        reason: e.to_string(),
    })?;
    let actual = sha256_hex(&bytes);
    if actual != file.sha256 {
        return Err(OcrError::ChecksumMismatch {
            path: path.display().to_string(),
            expected: file.sha256.to_string(),
            actual,
        });
    }
    // Write via a temp sibling + atomic rename so a concurrent reader can
    // never observe a partially-written model file.
    let tmp = path.with_extension("tmp");
    tokio::fs::write(&tmp, &bytes)
        .await
        .map_err(|e| OcrError::Download {
            url: file.url.to_string(),
            reason: format!("cache write failed: {e}"),
        })?;
    tokio::fs::rename(&tmp, &path)
        .await
        .map_err(|e| OcrError::Download {
            url: file.url.to_string(),
            reason: format!("cache rename failed: {e}"),
        })?;
    Ok(path)
}

/// Ensure all three artifacts are cached; returns (det, rec, dict) paths.
pub async fn ensure_all_cached() -> Result<(PathBuf, PathBuf, PathBuf), OcrError> {
    let dir = default_cache_dir();
    Ok((
        ensure_cached(&DET_MODEL, &dir).await?,
        ensure_cached(&REC_MODEL, &dir).await?,
        ensure_cached(&REC_DICT, &dir).await?,
    ))
}

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

    const HELLO_SHA256: &str = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824";

    #[test]
    fn sha256_hex_computes_known_digest() {
        assert_eq!(sha256_hex(b"hello"), HELLO_SHA256);
    }

    #[tokio::test]
    async fn ensure_cached_accepts_preseeded_valid_file() {
        let dir = tempfile::tempdir().unwrap();
        let file = ModelFile {
            filename: "hello.bin",
            // Bogus URL: the network must never be touched when the cache is valid.
            url: "https://invalid.invalid/hello.bin",
            sha256: HELLO_SHA256,
        };
        std::fs::write(dir.path().join("hello.bin"), b"hello").unwrap();
        let path = ensure_cached(&file, dir.path()).await.unwrap();
        assert_eq!(path, dir.path().join("hello.bin"));
    }

    #[tokio::test]
    async fn ensure_cached_deletes_corrupt_file_and_errors() {
        let dir = tempfile::tempdir().unwrap();
        let file = ModelFile {
            filename: "hello.bin",
            // Port 1 on loopback refuses the connection immediately -- no DNS
            // dependency, so the test stays hermetic and fast everywhere.
            url: "http://127.0.0.1:1/hello.bin",
            sha256: HELLO_SHA256,
        };
        std::fs::write(dir.path().join("hello.bin"), b"corrupted").unwrap();
        // Corrupt cache -> deleted -> re-download attempted -> Download error
        // (the connection is refused by construction).
        let err = ensure_cached(&file, dir.path()).await.unwrap_err();
        assert!(matches!(err, OcrError::Download { .. }), "got {err:?}");
        assert!(
            !dir.path().join("hello.bin").exists(),
            "corrupt file must be deleted"
        );
    }
}