regy 0.1.4

Private-by-default desktop agent for the Regy web interface
use std::{
    fs::{self, File},
    io::Read,
    path::{Path, PathBuf},
};

use serde::Deserialize;

use crate::domain::errors::{AgentError, AgentResult, ErrorCode};

const MAX_DISCOVERY_DEPTH: usize = 16;
const MAX_DISCOVERY_ENTRIES: usize = 4096;
const MAX_DISCOVERY_FILES: usize = 2048;
const MAX_DISCOVERY_BYTES: usize = 2 * 1024 * 1024;
const MAX_DISCOVERED_SKILLS: usize = 512;
const MAX_SKILL_MARKDOWN_BYTES: usize = 256 * 1024;

#[derive(Debug)]
pub(crate) struct PiSkillSource {
    pub(crate) canonical_path: PathBuf,
    pub(crate) name: String,
    pub(crate) description: String,
}

#[derive(Default)]
struct DiscoveryBudget {
    entries: usize,
    files: usize,
    bytes: usize,
    skills: usize,
}

#[derive(Deserialize)]
struct LenientSkillMetadata {
    name: String,
    description: String,
}

pub(crate) fn discover_pi_skills(
    root: &Path,
    accept_root_markdown: bool,
) -> AgentResult<Vec<PiSkillSource>> {
    let metadata = match fs::symlink_metadata(root) {
        Ok(metadata) => metadata,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(_) => return Err(filesystem_failed()),
    };
    if metadata.file_type().is_symlink() || !metadata.file_type().is_dir() {
        return Ok(Vec::new());
    }

    let mut discovered = Vec::new();
    let mut budget = DiscoveryBudget::default();
    let mut pending = vec![(root.to_path_buf(), 0usize, true)];
    while let Some((directory, depth, is_root)) = pending.pop() {
        if budget.entries >= MAX_DISCOVERY_ENTRIES
            || budget.files >= MAX_DISCOVERY_FILES
            || budget.bytes >= MAX_DISCOVERY_BYTES
            || budget.skills >= MAX_DISCOVERED_SKILLS
        {
            break;
        }

        if !is_root {
            append_skill_markdown(&directory.join("SKILL.md"), &mut budget, &mut discovered)?;
        }

        let entries = match fs::read_dir(&directory) {
            Ok(entries) => entries,
            Err(_) => continue,
        };
        let mut entries = entries.filter_map(Result::ok).collect::<Vec<_>>();
        entries.sort_by_key(|entry| entry.file_name());
        for entry in entries.into_iter().rev() {
            budget.entries += 1;
            if budget.entries > MAX_DISCOVERY_ENTRIES {
                break;
            }
            let file_type = match entry.file_type() {
                Ok(file_type) => file_type,
                Err(_) => continue,
            };
            if file_type.is_symlink() {
                continue;
            }
            let name = entry.file_name();
            let name = name.to_string_lossy();
            if file_type.is_dir() {
                if is_transient_directory(&name) || depth >= MAX_DISCOVERY_DEPTH {
                    continue;
                }
                pending.push((entry.path(), depth + 1, false));
            } else if file_type.is_file() {
                budget.files += 1;
                if is_root
                    && accept_root_markdown
                    && entry
                        .path()
                        .extension()
                        .is_some_and(|extension| extension == "md")
                {
                    append_skill_markdown(&entry.path(), &mut budget, &mut discovered)?;
                }
            }
        }
    }
    Ok(discovered)
}

fn append_skill_markdown(
    path: &Path,
    budget: &mut DiscoveryBudget,
    discovered: &mut Vec<PiSkillSource>,
) -> AgentResult<()> {
    if budget.skills >= MAX_DISCOVERED_SKILLS || budget.bytes >= MAX_DISCOVERY_BYTES {
        return Ok(());
    }
    let metadata = match fs::symlink_metadata(path) {
        Ok(metadata) => metadata,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
        Err(_) => return Ok(()),
    };
    if metadata.file_type().is_symlink()
        || !metadata.file_type().is_file()
        || metadata.len() > MAX_SKILL_MARKDOWN_BYTES as u64
    {
        return Ok(());
    }
    let remaining = MAX_DISCOVERY_BYTES.saturating_sub(budget.bytes);
    let limit = remaining.min(MAX_SKILL_MARKDOWN_BYTES);
    let mut bytes = Vec::new();
    File::open(path)
        .and_then(|file| file.take((limit + 1) as u64).read_to_end(&mut bytes))
        .map_err(|_| filesystem_failed())?;
    if bytes.len() > limit {
        return Ok(());
    }
    budget.bytes += bytes.len();
    let Some((name, description)) = parse_lenient_skill_metadata(&bytes) else {
        return Ok(());
    };
    let canonical_path = match path.canonicalize() {
        Ok(path) => path,
        Err(_) => return Ok(()),
    };
    budget.skills += 1;
    discovered.push(PiSkillSource {
        canonical_path,
        name,
        description,
    });
    Ok(())
}

pub(crate) fn parse_lenient_skill_metadata(bytes: &[u8]) -> Option<(String, String)> {
    let yaml = frontmatter(bytes)?;
    let mut metadata = serde_yaml_ng::from_str::<LenientSkillMetadata>(&yaml).ok()?;
    metadata.name = metadata.name.trim().to_owned();
    metadata.description = metadata.description.trim().to_owned();
    (!metadata.name.is_empty() && !metadata.description.is_empty())
        .then_some((metadata.name, metadata.description))
}

pub(crate) fn frontmatter(bytes: &[u8]) -> Option<String> {
    let contents = std::str::from_utf8(bytes).ok()?;
    let normalized = contents.replace("\r\n", "\n");
    let mut lines = normalized.split_inclusive('\n');
    (lines.next().map(str::trim_end) == Some("---")).then_some(())?;
    let mut yaml = String::new();
    for line in lines {
        if line.trim_end() == "---" {
            return Some(yaml);
        }
        yaml.push_str(line);
    }
    None
}

fn is_transient_directory(name: &str) -> bool {
    name.starts_with(".regy-stage-")
        || name.starts_with(".regy-quarantine-")
        || name.starts_with(".regy-failed-")
}

fn filesystem_failed() -> AgentError {
    AgentError::new(
        ErrorCode::SkillFilesystemFailed,
        "skill filesystem discovery failed",
    )
}