PicoBot/build.rs

use std::env;
use std::fs;
use std::io::Write;
use std::path::Path;
use std::process::Command;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    build_webui(Path::new(&out_dir));

    println!("cargo:rerun-if-changed=resources/skills");
    let skills_dir = Path::new("resources/skills");
    let skills_out_dir = Path::new(&out_dir).join("skills");
    fs::create_dir_all(&skills_out_dir).unwrap();

    println!("cargo:rerun-if-changed=resources/agents");
    let agents_dir = Path::new("resources/agents");
    let agents_out_dir = Path::new(&out_dir).join("agents");
    fs::create_dir_all(&agents_out_dir).unwrap();
    let mut agents = Vec::new();
    if let Ok(entries) = fs::read_dir(agents_dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("md") {
                continue;
            }
            let agent_name = path
                .file_stem()
                .unwrap()
                .to_str()
                .unwrap()
                .to_string();
            fs::copy(&path, agents_out_dir.join(format!("{agent_name}.md"))).unwrap();
            agents.push(agent_name);
        }
    }
    agents.sort();

    let mut skills = Vec::new();

    if let Ok(entries) = fs::read_dir(skills_dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if !path.is_dir() {
                continue;
            }

            let skill_name = path.file_name().unwrap().to_str().unwrap().to_string();
            let compressed_path = skills_out_dir.join(format!("{}.tar.zst", skill_name));

            let compressed = compress_skill_dir(&path);
            fs::write(&compressed_path, &compressed).unwrap();

            skills.push(skill_name);
        }
    }

    let mut code = String::from(
        r#"pub struct EmbeddedSkill {
    pub name: &'static str,
    pub data: &'static [u8],
}

pub static EMBEDDED_SKILLS: &[EmbeddedSkill] = &[
"#,
    );

    for name in &skills {
        let file_path = skills_out_dir
            .join(format!("{}.tar.zst", name))
            .to_string_lossy()
            .to_string();
        code.push_str(&format!(
            "    EmbeddedSkill {{ name: \"{}\", data: include_bytes!(\"{}\") }},\n",
            name, file_path
        ));
    }

    code.push_str("];\n");

    let generated_path = Path::new(&out_dir).join("embedded_skills.rs");
    let mut f = fs::File::create(&generated_path).unwrap();
    f.write_all(code.as_bytes()).unwrap();

    let mut agent_code = String::from(
        r#"pub struct EmbeddedAgent {
    pub name: &'static str,
    pub content: &'static str,
}

pub static EMBEDDED_AGENTS: &[EmbeddedAgent] = &[
"#,
    );
    for name in &agents {
        let file_path = agents_out_dir
            .join(format!("{name}.md"))
            .to_string_lossy()
            .to_string();
        agent_code.push_str(&format!(
            "    EmbeddedAgent {{ name: \"{name}\", content: include_str!(\"{file_path}\") }},\n",
        ));
    }
    agent_code.push_str("];\n");
    let agent_path = Path::new(&out_dir).join("embedded_agents.rs");
    let mut f = fs::File::create(&agent_path).unwrap();
    f.write_all(agent_code.as_bytes()).unwrap();
}

fn build_webui(out_dir: &Path) {
    for path in [
        "webui/src",
        "webui/index.html",
        "webui/jsconfig.json",
        "webui/package.json",
        "webui/package-lock.json",
        "webui/svelte.config.js",
        "webui/vite.config.js",
        "webui/public",
    ] {
        println!("cargo:rerun-if-changed={path}");
    }

    let webui_dir = Path::new("webui");
    let lockfile = webui_dir.join("package-lock.json");
    let dependency_stamp = webui_dir
        .join("node_modules")
        .join(".picobot-package-lock.json");
    let lockfile_contents = fs::read(&lockfile).expect("failed to read webui/package-lock.json");
    let dependencies_current = fs::read(&dependency_stamp)
        .is_ok_and(|stamp| stamp == lockfile_contents)
        && webui_dir.join("node_modules/.bin/vite").is_file();

    if !dependencies_current {
        run_npm(webui_dir, &["ci", "--no-audit", "--no-fund"], None);
        fs::write(&dependency_stamp, &lockfile_contents)
            .expect("failed to write WebUI dependency stamp");
    }

    let webui_out_dir = out_dir.join("webui");
    run_npm(webui_dir, &["run", "build"], Some(&webui_out_dir));
}

fn run_npm(webui_dir: &Path, args: &[&str], output_dir: Option<&Path>) {
    let mut command = Command::new("npm");
    command.args(args).current_dir(webui_dir);
    if let Some(output_dir) = output_dir {
        command.env("PICOBOT_WEBUI_OUT_DIR", output_dir);
    }

    let output = command.output().unwrap_or_else(|error| {
        panic!("failed to start npm for the WebUI build ({error}); install Node.js 20+ and npm")
    });
    if !output.status.success() {
        panic!(
            "WebUI command `npm {}` failed\nstdout:\n{}\nstderr:\n{}",
            args.join(" "),
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
    }
}

fn compress_skill_dir(dir: &Path) -> Vec<u8> {
    let mut buf = Vec::new();
    let mut builder = tar::Builder::new(&mut buf);
    builder.follow_symlinks(false);
    builder
        .append_dir_all(dir.file_name().unwrap().to_str().unwrap(), dir)
        .unwrap();
    drop(builder);

    zstd::encode_all(buf.as_slice(), 3).unwrap()
}