pdf-ocr 1.0.0

//! Model loading, downloading, and caching for PaddleOCR.

// Copyright (c) 2026 Innovation Trigger B.V.
//
// PDFluent is available under two licences, at your option: the GNU AGPLv3, or
// the PDFluent Commercial Licence. See the LICENSE file in this repository --
// that file travels with the copy you received, which a URL does not.

use std::io::Write;
use std::path::{Path, PathBuf};

use ort::session::Session;

use super::dictionary::Dictionary;

/// Error type for model operations.
#[derive(Debug)]
pub enum ModelError {
    /// File system I/O error while reading or writing model files.
    Io(std::io::Error),
    /// HTTP download of the model file failed; contains the error description.
    Download(String),
    /// Creating the ONNX runtime session failed; contains the error description.
    Session(String),
    /// Loading the character dictionary failed; contains the error description.
    Dictionary(String),
    /// Model files are absent and `ModelSource::LocalOnly` forbids fetching them.
    ModelsMissing {
        /// The files that are not present.
        missing: Vec<PathBuf>,
        /// Where they were looked for.
        model_dir: PathBuf,
    },
    /// A fetch was requested for a file with no pinned digest.
    UnpinnedFile(String),
    /// Bytes arrived but did not match the pinned digest. Nothing was written.
    Integrity(String),
}

impl std::fmt::Display for ModelError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(e) => write!(f, "model I/O error: {e}"),
            Self::Download(msg) => write!(f, "model download error: {msg}"),
            Self::Session(msg) => write!(f, "session creation error: {msg}"),
            Self::Dictionary(msg) => write!(f, "dictionary error: {msg}"),
            Self::ModelsMissing { missing, model_dir } => write!(
                f,
                "OCR model files are missing from {}: {}. \
                 Nothing was downloaded, by design: place the weights there yourself, \
                 or set PaddleOcrConfig::model_source to ModelSource::Verified with a \
                 pinned sha256 for each file. Reference source: {}",
                model_dir.display(),
                missing
                    .iter()
                    .map(|p| p
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy()
                        .into_owned())
                    .collect::<Vec<_>>()
                    .join(", "),
                REFERENCE_MODEL_BASE_URL
            ),
            Self::UnpinnedFile(path) => write!(
                f,
                "refusing to download {path}: no sha256 pinned for it. Add one to \
                 ModelSource::Verified::digests -- a file that cannot be verified \
                 must not be fetched"
            ),
            Self::Integrity(msg) => write!(f, "{msg}"),
        }
    }
}

impl std::error::Error for ModelError {}

impl From<std::io::Error> for ModelError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

/// Loaded ONNX sessions and dictionary for PaddleOCR inference.
pub struct ModelSessions {
    /// ONNX session for the DBNet text detection model.
    pub detection: Session,
    /// ONNX session for the SVTR text recognition model.
    pub recognition: Session,
    /// Optional ONNX session for the angle classifier (180° rotation detection).
    pub classifier: Option<Session>,
    /// Character dictionary mapping label indices to Unicode strings.
    pub dictionary: Dictionary,
}

/// Detection model variant.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DetectionModel {
    /// PP-OCRv3 (2.3 MB) — lightweight, good for most documents.
    V3,
    /// PP-OCRv5 (84 MB) — best quality.
    V5,
}

/// Recognition language.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Language {
    /// English language recognition model (`en`).
    English,
    /// Latin-script language recognition model (`latin`).
    Latin,
    /// Simplified Chinese recognition model (`ch`).
    Chinese,
    /// Japanese recognition model (`japan`).
    Japanese,
    /// Korean recognition model (`korean`).
    Korean,
    /// Arabic-script recognition model (`arabic`).
    Arabic,
}

impl Language {
    /// Language code string.
    pub fn code(self) -> &'static str {
        match self {
            Self::English => "en",
            Self::Latin => "latin",
            Self::Chinese => "ch",
            Self::Japanese => "japan",
            Self::Korean => "korean",
            Self::Arabic => "arabic",
        }
    }

    /// HuggingFace directory path for this language (used for downloads).
    fn hf_dir(self) -> &'static str {
        match self {
            Self::English => "languages/english",
            Self::Latin => "languages/latin",
            Self::Chinese => "languages/chinese",
            Self::Japanese => "languages/japan",
            Self::Korean => "languages/korean",
            Self::Arabic => "languages/arabic",
        }
    }

    /// Local directory name for cached models.
    fn local_dir(self) -> &'static str {
        match self {
            Self::English => "english",
            Self::Latin => "latin",
            Self::Chinese => "chinese",
            Self::Japanese => "japan",
            Self::Korean => "korean",
            Self::Arabic => "arabic",
        }
    }
}

/// Configuration for the PaddleOCR engine.
pub struct PaddleOcrConfig {
    /// Directory containing ONNX model files.
    pub model_dir: PathBuf,
    /// Detection model variant.
    pub detection_model: DetectionModel,
    /// Languages for recognition (determines which rec model + dict to load).
    pub languages: Vec<Language>,
    /// Whether to use the angle classifier (for rotated text).
    pub use_angle_classifier: bool,
    /// Maximum side length for detection input (default 960).
    pub max_side_len: u32,
    /// Detection threshold (default 0.3).
    pub det_threshold: f32,
    /// Box threshold (default 0.6).
    pub box_threshold: f32,
    /// Recognition batch size (default 8).
    pub rec_batch_size: usize,
    /// Number of ONNX Runtime intra-op threads (default 4).
    pub num_threads: usize,
    /// Where weights may be fetched from. Defaults to `LocalOnly` — no network.
    pub model_source: ModelSource,
}

impl Default for PaddleOcrConfig {
    fn default() -> Self {
        Self {
            model_dir: default_cache_dir(),
            detection_model: DetectionModel::V3,
            languages: vec![Language::Latin],
            use_angle_classifier: false,
            max_side_len: 960,
            det_threshold: 0.3,
            box_threshold: 0.6,
            rec_batch_size: 8,
            num_threads: 4,
            model_source: ModelSource::LocalOnly,
        }
    }
}

fn default_cache_dir() -> PathBuf {
    dirs_next::cache_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("xfa")
        .join("ocr-models")
}

/// Get the local path for the detection model.
pub fn detection_model_path(config: &PaddleOcrConfig) -> PathBuf {
    let dir = match config.detection_model {
        DetectionModel::V3 => "det_v3",
        DetectionModel::V5 => "det_v5",
    };
    config.model_dir.join(dir).join("det.onnx")
}

/// Get the local path for the recognition model (first language).
pub fn recognition_model_path(config: &PaddleOcrConfig) -> PathBuf {
    let lang = config.languages.first().copied().unwrap_or(Language::Latin);
    config.model_dir.join(lang.local_dir()).join("rec.onnx")
}

/// Get the local path for the dictionary file (first language).
pub fn dictionary_path(config: &PaddleOcrConfig) -> PathBuf {
    let lang = config.languages.first().copied().unwrap_or(Language::Latin);
    config.model_dir.join(lang.local_dir()).join("dict.txt")
}

/// Get the local path for the angle classifier model.
pub fn classifier_model_path(config: &PaddleOcrConfig) -> PathBuf {
    config.model_dir.join("cls").join("cls.onnx")
}

/// Check if required models are available locally.
pub fn models_available(config: &PaddleOcrConfig) -> bool {
    let det = detection_model_path(config).exists();
    let rec = recognition_model_path(config).exists();
    let dict = dictionary_path(config).exists();
    let cls = if config.use_angle_classifier {
        classifier_model_path(config).exists()
    } else {
        true
    };
    det && rec && dict && cls
}

/// Where model weights may come from.
///
/// # Why the default fetches nothing
///
/// This used to download up to ~84 MB of ONNX weights from a third-party public
/// repository the first time an engine was constructed, and write them straight to
/// the cache without checking a single byte. Whoever served those bytes decided
/// what our inference engine executed, and the caller was never asked. For an SDK
/// sold on privacy and on working offline, a silent network fetch is the wrong
/// default even before the integrity question.
///
/// So the default is `LocalOnly`: point `model_dir` at weights you already have.
/// Fetching is a decision the caller makes explicitly, and it is always verified.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum ModelSource {
    /// Never touch the network. Missing files are an error naming what is missing.
    #[default]
    LocalOnly,
    /// Fetch missing files from `base_url`, verifying each against a pinned digest.
    ///
    /// Keys are repository-relative paths exactly as they appear in the URL, e.g.
    /// `"detection/v3/det.onnx"`. A file with no entry is **not** downloaded: an
    /// unpinned file cannot be verified, and downloading what you cannot verify is
    /// the thing this type exists to prevent.
    Verified {
        /// Base URL, without a trailing slash.
        base_url: String,
        /// Repository-relative path to hex SHA-256.
        digests: std::collections::BTreeMap<String, String>,
    },
}

/// Upper bound on a single model file, as a guard against a hostile or broken
/// server streaming without end. The largest weight we know of is ~84 MB.
const MAX_MODEL_BYTES: usize = 256 * 1024 * 1024;

/// The reference URL for PaddleOCR ONNX weights.
///
/// Provided for convenience only. It is a third-party repository we do not
/// control, which is exactly why using it requires pinning digests yourself.
pub const REFERENCE_MODEL_BASE_URL: &str =
    "https://huggingface.co/monkt/paddleocr-onnx/resolve/main";

/// Every file this configuration needs, as (repository-relative path, local path).
fn required_files(config: &PaddleOcrConfig) -> Vec<(String, PathBuf)> {
    let det_hf_dir = match config.detection_model {
        DetectionModel::V3 => "detection/v3",
        DetectionModel::V5 => "detection/v5",
    };
    let lang = config.languages.first().copied().unwrap_or(Language::Latin);

    let mut files = vec![
        (
            format!("{det_hf_dir}/det.onnx"),
            detection_model_path(config),
        ),
        (
            format!("{}/rec.onnx", lang.hf_dir()),
            recognition_model_path(config),
        ),
        (
            format!("{}/dict.txt", lang.hf_dir()),
            dictionary_path(config),
        ),
    ];
    if config.use_angle_classifier {
        files.push((
            "preprocessing/cls.onnx".to_string(),
            classifier_model_path(config),
        ));
    }
    files
}

/// Make sure every required model file is present, fetching only if the caller
/// asked for that and only against a pinned digest.
pub fn ensure_models(config: &PaddleOcrConfig) -> Result<(), ModelError> {
    let missing: Vec<(String, PathBuf)> = required_files(config)
        .into_iter()
        .filter(|(_, local)| !local.exists())
        .collect();

    if missing.is_empty() {
        return Ok(());
    }

    match &config.model_source {
        ModelSource::LocalOnly => Err(ModelError::ModelsMissing {
            missing: missing.iter().map(|(_, p)| p.clone()).collect(),
            model_dir: config.model_dir.clone(),
        }),
        ModelSource::Verified { base_url, digests } => {
            for (hf_path, local_path) in missing {
                let expected = digests
                    .get(&hf_path)
                    .ok_or_else(|| ModelError::UnpinnedFile(hf_path.clone()))?;
                fetch_verified(base_url, &hf_path, expected, &local_path)?;
            }
            Ok(())
        }
    }
}

/// Download one file into memory, verify it, and only then write it.
///
/// The order matters. Streaming to disk and checking afterwards leaves unverified
/// bytes on the filesystem, where a crash, a concurrent process, or a caller that
/// ignores the error can still pick them up and load them.
fn fetch_verified(
    base_url: &str,
    hf_path: &str,
    expected_sha256: &str,
    local_path: &Path,
) -> Result<(), ModelError> {
    let url = format!("{}/{}", base_url.trim_end_matches('/'), hf_path);

    let agent = ureq::Agent::new_with_defaults();
    let response = agent
        .get(&url)
        .call()
        .map_err(|e| ModelError::Download(format!("{url}: {e}")))?;

    let mut reader = response.into_body().into_reader();
    let mut bytes: Vec<u8> = Vec::new();
    let mut buf = vec![0u8; 64 * 1024];
    loop {
        let n = std::io::Read::read(&mut reader, &mut buf)
            .map_err(|e| ModelError::Download(format!("{url}: read error: {e}")))?;
        if n == 0 {
            break;
        }
        if bytes.len() + n > MAX_MODEL_BYTES {
            return Err(ModelError::Download(format!(
                "{url}: exceeds {MAX_MODEL_BYTES} bytes; refusing to continue"
            )));
        }
        bytes.extend_from_slice(&buf[..n]);
    }

    crate::integrity::verify_sha256(&bytes, expected_sha256)
        .map_err(|e| ModelError::Integrity(format!("{url}: {e}")))?;

    if let Some(parent) = local_path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    // Write beside the target and rename, so a reader never observes a partial
    // file even though the bytes are already verified.
    let tmp = local_path.with_extension("part");
    {
        let mut file = std::fs::File::create(&tmp)?;
        file.write_all(&bytes)?;
        file.sync_all()?;
    }
    std::fs::rename(&tmp, local_path)?;
    Ok(())
}

/// Load ONNX sessions for all required models.
///
/// **No test reaches this, deliberately, and that follows the scope decision
/// rather than a gap in the suite.**
///
/// BESLUIT 19-08-2026: we do not develop OCR. We facilitate three routes to
/// somebody else's -- cloud through the `OcrEngine` trait, PaddleOCR behind the
/// `paddle` feature, Tesseract behind `tesseract` -- and the facade wires none
/// of them, so an ordinary `pdfluent` dependency pulls no system libraries and
/// downloads no models.
///
/// This function is inside the `paddle` route. It drives a third-party ONNX
/// model that is fetched once from HuggingFace and never committed. Testing it
/// would assert that ONNX Runtime returns tensors for weights we did not train
/// and do not maintain -- which is not our behaviour to guard.
///
/// What is ours is the seam: the `OcrEngine` trait and `make_searchable`, which
/// writes words and `bbox_px` back as an invisible text layer. That is where a
/// test earns its keep.
pub fn load_sessions(config: &PaddleOcrConfig) -> Result<ModelSessions, ModelError> {
    let detection = create_session(&detection_model_path(config), config.num_threads)?;
    let recognition = create_session(&recognition_model_path(config), config.num_threads)?;

    let classifier = if config.use_angle_classifier {
        Some(create_session(
            &classifier_model_path(config),
            config.num_threads,
        )?)
    } else {
        None
    };

    let dictionary = Dictionary::from_file(&dictionary_path(config))
        .map_err(|e| ModelError::Dictionary(e.to_string()))?;

    Ok(ModelSessions {
        detection,
        recognition,
        classifier,
        dictionary,
    })
}

fn create_session(model_path: &Path, num_threads: usize) -> Result<Session, ModelError> {
    let session = Session::builder()
        .map_err(|e| ModelError::Session(e.to_string()))?
        .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)
        .map_err(|e| ModelError::Session(e.to_string()))?
        .with_intra_threads(num_threads)
        .map_err(|e| ModelError::Session(e.to_string()))?
        .commit_from_file(model_path)
        .map_err(|e| ModelError::Session(format!("{}: {e}", model_path.display())))?;
    Ok(session)
}

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

    #[test]
    fn default_config_is_sensible() {
        let config = PaddleOcrConfig::default();
        assert_eq!(config.max_side_len, 960);
        assert!((config.det_threshold - 0.3).abs() < f32::EPSILON);
        assert_eq!(config.rec_batch_size, 8);
        assert_eq!(config.num_threads, 4);
        assert!(!config.use_angle_classifier);
    }

    #[test]
    fn language_code_mapping() {
        assert_eq!(Language::English.code(), "en");
        assert_eq!(Language::Latin.code(), "latin");
        assert_eq!(Language::Chinese.code(), "ch");
        assert_eq!(Language::Japanese.code(), "japan");
        assert_eq!(Language::Korean.code(), "korean");
        assert_eq!(Language::Arabic.code(), "arabic");
    }

    #[test]
    fn model_path_construction_v3() {
        let config = PaddleOcrConfig {
            model_dir: PathBuf::from("/tmp/models"),
            detection_model: DetectionModel::V3,
            languages: vec![Language::English],
            ..Default::default()
        };
        let det_path = detection_model_path(&config);
        assert!(det_path.to_str().unwrap().contains("det_v3"));
        assert!(det_path.to_str().unwrap().ends_with("det.onnx"));
    }

    #[test]
    fn model_path_construction_v5() {
        let config = PaddleOcrConfig {
            model_dir: PathBuf::from("/tmp/models"),
            detection_model: DetectionModel::V5,
            languages: vec![Language::Latin],
            ..Default::default()
        };
        let det_path = detection_model_path(&config);
        assert!(det_path.to_str().unwrap().contains("det_v5"));

        let rec_path = recognition_model_path(&config);
        assert!(rec_path.to_str().unwrap().contains("latin"));
        assert!(rec_path.to_str().unwrap().ends_with("rec.onnx"));

        let dict_path = dictionary_path(&config);
        assert!(dict_path.to_str().unwrap().contains("latin"));
        assert!(dict_path.to_str().unwrap().ends_with("dict.txt"));
    }

    #[test]
    fn models_not_available_when_missing() {
        let config = PaddleOcrConfig {
            model_dir: PathBuf::from("/nonexistent/path"),
            ..Default::default()
        };
        assert!(!models_available(&config));
    }
}