siteforge 0.1.4

Archive websites into AI-readable local knowledge archives
Documentation
use std::path::Path;
use std::process::Command;

use crate::config::Config;
use crate::errors::{Result, SiteforgeError};
use crate::parse::AssetRef;

#[derive(Debug, Clone)]
pub struct OcrOutput {
    pub text: String,
    pub confidence: f32,
}

pub trait OcrBackend: Send + Sync {
    fn name(&self) -> &str;
    fn is_available(&self) -> Result<()>;
    fn extract(&self, image_path: &Path) -> Result<OcrOutput>;
}

#[derive(Debug, Clone)]
pub struct TesseractOcr {
    command_template: String,
}

impl TesseractOcr {
    pub fn new(command_template: impl Into<String>) -> Self {
        Self {
            command_template: command_template.into(),
        }
    }

    fn command_parts(&self, input: Option<&Path>) -> Result<Vec<String>> {
        let rendered = match input {
            Some(path) => self
                .command_template
                .replace("{input}", &path.to_string_lossy()),
            None => self.command_template.clone(),
        };
        shell_words::split(&rendered).map_err(|err| SiteforgeError::message(err.to_string()))
    }
}

impl OcrBackend for TesseractOcr {
    fn name(&self) -> &str {
        "tesseract"
    }

    fn is_available(&self) -> Result<()> {
        let parts = self.command_parts(None)?;
        let executable = parts
            .first()
            .ok_or_else(|| SiteforgeError::OcrUnavailable("empty OCR command".to_string()))?;
        let status = Command::new(executable).arg("--version").status();
        match status {
            Ok(status) if status.success() => Ok(()),
            Ok(_) => Err(SiteforgeError::OcrUnavailable(format!(
                "{executable} --version returned a non-zero status"
            ))),
            Err(err) => Err(SiteforgeError::OcrUnavailable(format!(
                "{executable} is not available: {err}"
            ))),
        }
    }

    fn extract(&self, image_path: &Path) -> Result<OcrOutput> {
        self.is_available()?;
        let mut parts = self.command_parts(Some(image_path))?;
        if parts.is_empty() {
            return Err(SiteforgeError::OcrUnavailable(
                "empty OCR command".to_string(),
            ));
        }
        if !self.command_template.contains("{input}") {
            parts.push(image_path.to_string_lossy().to_string());
            parts.push("stdout".to_string());
        }
        let executable = parts.remove(0);
        let output = Command::new(executable).args(parts).output()?;
        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(SiteforgeError::message(format!(
                "OCR command failed: {stderr}"
            )));
        }
        Ok(OcrOutput {
            text: String::from_utf8_lossy(&output.stdout).trim().to_string(),
            confidence: 0.55,
        })
    }
}

pub fn backend_from_config(config: &Config) -> Option<TesseractOcr> {
    config
        .ocr_enabled
        .then(|| TesseractOcr::new(config.ocr_backend_command.clone()))
}

pub fn likely_contains_code(asset: &AssetRef) -> bool {
    let haystack = [
        asset.url.as_str(),
        asset.alt_text.as_deref().unwrap_or_default(),
        asset.caption.as_deref().unwrap_or_default(),
    ]
    .join(" ")
    .to_ascii_lowercase();
    let keyword_hit = [
        "code",
        "script",
        "snippet",
        "source",
        "screenshot",
        "terminal",
        "console",
        "error",
        "stacktrace",
        "unity",
        "shader",
    ]
    .iter()
    .any(|keyword| haystack.contains(keyword));
    let extension_hit = [".png", ".jpg", ".jpeg", ".webp"]
        .iter()
        .any(|ext| haystack.contains(ext));
    keyword_hit && extension_hit
}

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

    #[test]
    fn unavailable_ocr_is_clear() {
        let backend = TesseractOcr::new("definitely-not-installed-siteforge-ocr {input} stdout");
        let err = backend.is_available().unwrap_err().to_string();
        assert!(err.contains("OCR unavailable"));
    }
}