siteforge 0.1.4

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

use url::Url;

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

#[derive(Debug, Clone)]
pub struct BrowserRenderOptions {
    pub command: Option<String>,
    pub profile: Option<PathBuf>,
    pub wait_ms: u64,
    pub user_agent: String,
}

impl BrowserRenderOptions {
    pub fn from_config(config: &Config) -> Self {
        Self {
            command: config.browser_command.clone(),
            profile: config.browser_profile.clone(),
            wait_ms: config.render_wait_ms,
            user_agent: config.user_agent.clone(),
        }
    }
}

pub async fn render_js_dom(options: BrowserRenderOptions, url: Url) -> Result<String> {
    tokio::task::spawn_blocking(move || render_js_dom_blocking(&options, &url))
        .await
        .map_err(|err| SiteforgeError::message(format!("browser render task failed: {err}")))?
}

fn render_js_dom_blocking(options: &BrowserRenderOptions, url: &Url) -> Result<String> {
    if let Some(command) = options.command.as_deref() {
        return run_command(command, options, url);
    }

    let candidates = [
        "chromium",
        "chromium-browser",
        "google-chrome",
        "google-chrome-stable",
    ];
    let mut errors = Vec::new();
    for candidate in candidates {
        match run_browser(candidate, &[], options, url) {
            Ok(html) => return Ok(html),
            Err(err) => errors.push(format!("{candidate}: {err}")),
        }
    }
    Err(SiteforgeError::message(format!(
        "no working browser renderer found; set browser_command. Tried: {}",
        errors.join("; ")
    )))
}

fn run_command(command: &str, options: &BrowserRenderOptions, url: &Url) -> Result<String> {
    let rendered = command
        .replace("{url}", url.as_str())
        .replace("{wait_ms}", &options.wait_ms.to_string())
        .replace("{user_agent}", &options.user_agent);
    let parts = shell_words::split(&rendered)
        .map_err(|err| SiteforgeError::message(format!("invalid browser command: {err}")))?;
    let (program, args) = parts
        .split_first()
        .ok_or_else(|| SiteforgeError::message("empty browser command"))?;

    if command.contains("{url}") {
        run_raw_command(program, args)
    } else {
        run_browser(program, args, options, url)
    }
}

fn run_browser(
    program: &str,
    prefix_args: &[String],
    options: &BrowserRenderOptions,
    url: &Url,
) -> Result<String> {
    let mut args = prefix_args.to_vec();
    args.extend([
        "--headless".to_string(),
        "--disable-gpu".to_string(),
        "--no-sandbox".to_string(),
        "--disable-dev-shm-usage".to_string(),
        "--dump-dom".to_string(),
        format!("--virtual-time-budget={}", options.wait_ms),
        format!("--user-agent={}", options.user_agent),
    ]);
    if let Some(profile) = options.profile.as_deref() {
        args.push(format!("--user-data-dir={}", profile.display()));
    }
    args.push(url.to_string());
    run_raw_command(program, &args)
}

fn run_raw_command(program: &str, args: &[String]) -> Result<String> {
    let output = Command::new(program).args(args).output().map_err(|err| {
        SiteforgeError::message(format!("failed to start browser renderer {program}: {err}"))
    })?;
    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(SiteforgeError::message(format!(
            "browser renderer exited with {}: {}",
            output.status, stderr
        )));
    }
    let html = String::from_utf8_lossy(&output.stdout).trim().to_string();
    if html.is_empty() {
        return Err(SiteforgeError::message(
            "browser renderer returned empty HTML".to_string(),
        ));
    }
    Ok(html)
}

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

    #[test]
    fn custom_command_template_replaces_values() {
        let options = BrowserRenderOptions {
            command: Some("printf '<html>{url}:{wait_ms}</html>'".to_string()),
            profile: None,
            wait_ms: 42,
            user_agent: "agent".to_string(),
        };
        let url = Url::parse("https://example.test/").unwrap();
        let html = render_js_dom_blocking(&options, &url).unwrap();
        assert!(html.contains("https://example.test/:42"));
    }
}