PicoBot/src/skills/mod.rs

use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::SystemTime;

/// Skill definition
#[derive(Debug, Clone)]
pub struct Skill {
    pub name: String,
    pub description: String,
    pub content: String,
    pub always: bool,
    pub path: Option<PathBuf>,
}

#[derive(Default)]
struct SkillMarkdownMeta {
    name: Option<String>,
    description: Option<String>,
    always: Option<bool>,
}

#[derive(Clone)]
struct SkillsState {
    loaded_skills: Vec<Skill>,
    last_picobot_mtime: Option<SystemTime>,
    last_agent_mtime: Option<SystemTime>,
    last_workspace_mtime: Option<SystemTime>,
    last_load_time: SystemTime,
}

impl Default for SkillsState {
    fn default() -> Self {
        Self {
            loaded_skills: Vec::new(),
            last_picobot_mtime: None,
            last_agent_mtime: None,
            last_workspace_mtime: None,
            last_load_time: SystemTime::now(),
        }
    }
}

/// Skills loader - loads skills from multiple directories
#[derive(Clone)]
pub struct SkillsLoader {
    picobot_skills_dir: PathBuf,
    agent_skills_dir: PathBuf,
    workspace_skills_dir: Option<PathBuf>,
    state: Arc<Mutex<SkillsState>>,
}

impl SkillsLoader {
    /// Create a new loader with default paths
    pub fn new() -> Self {
        let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
        Self {
            picobot_skills_dir: home.join(".picobot/skills"),
            agent_skills_dir: home.join(".agent/skills"),
            workspace_skills_dir: None,
            state: Arc::new(Mutex::new(SkillsState::default())),
        }
    }

    #[cfg(test)]
    pub(crate) fn new_for_testing(picobot_dir: PathBuf, agent_dir: PathBuf) -> Self {
        Self {
            picobot_skills_dir: picobot_dir,
            agent_skills_dir: agent_dir,
            workspace_skills_dir: None,
            state: Arc::new(Mutex::new(SkillsState::default())),
        }
    }

    /// Set the workspace skills directory (./skills under workspace root)
    pub fn set_workspace_skills_dir(&mut self, workspace_path: PathBuf) {
        self.workspace_skills_dir = Some(workspace_path.join("skills"));
    }

    /// Load all skills from both directories and record modification times
    pub fn load_skills(&self) {
        let mut state = self.state.lock().unwrap();
        state.loaded_skills.clear();

        // Ensure ~/.picobot/skills directory exists
        if !self.picobot_skills_dir.exists() {
            if let Err(e) = std::fs::create_dir_all(&self.picobot_skills_dir) {
                tracing::warn!(dir = %self.picobot_skills_dir.display(), error = %e, "Failed to create skills directory");
            } else {
                tracing::info!(dir = %self.picobot_skills_dir.display(), "Created skills directory");
            }
        }

        // Load from ~/.picobot/skills
        if self.picobot_skills_dir.exists() {
            let loaded = self.load_skills_from_dir(&self.picobot_skills_dir);
            tracing::debug!(
                dir = %self.picobot_skills_dir.display(),
                count = loaded.len(),
                "Loaded skills from picobot directory"
            );
            state.loaded_skills.extend(loaded);
            state.last_picobot_mtime = Self::get_dir_mtime(&self.picobot_skills_dir);
        }

        // Load from ~/.agent/skills
        if self.agent_skills_dir.exists() {
            let loaded = self.load_skills_from_dir(&self.agent_skills_dir);
            tracing::debug!(
                dir = %self.agent_skills_dir.display(),
                count = loaded.len(),
                "Loaded skills from agent directory"
            );
            state.loaded_skills.extend(loaded);
            state.last_agent_mtime = Self::get_dir_mtime(&self.agent_skills_dir);
        }

        // Load from workspace ./skills (if set)
        if let Some(ref ws_dir) = self.workspace_skills_dir {
            if ws_dir.exists() {
                let loaded = self.load_skills_from_dir(ws_dir);
                tracing::debug!(
                    dir = %ws_dir.display(),
                    count = loaded.len(),
                    "Loaded skills from workspace directory"
                );
                state.loaded_skills.extend(loaded);
                state.last_workspace_mtime = Self::get_dir_mtime(ws_dir);
            }
        }

        state.last_load_time = SystemTime::now();

        if state.loaded_skills.is_empty() {
            tracing::debug!("No skills found in any skills directory");
        } else {
            tracing::info!(count = state.loaded_skills.len(), "Loaded {} skills total", state.loaded_skills.len());
        }
    }

    /// Check if skills directories have been modified since last load
    fn has_changed(&self) -> bool {
        let state = self.state.lock().unwrap();
        let picobot_changed = if self.picobot_skills_dir.exists() {
            let current_mtime = Self::get_dir_mtime(&self.picobot_skills_dir);
            current_mtime != state.last_picobot_mtime
        } else {
            false
        };

        let agent_changed = if self.agent_skills_dir.exists() {
            let current_mtime = Self::get_dir_mtime(&self.agent_skills_dir);
            current_mtime != state.last_agent_mtime
        } else {
            false
        };

        let workspace_changed = if let Some(ref ws_dir) = self.workspace_skills_dir {
            if ws_dir.exists() {
                let current_mtime = Self::get_dir_mtime(ws_dir);
                current_mtime != state.last_workspace_mtime
            } else {
                false
            }
        } else {
            false
        };

        picobot_changed || agent_changed || workspace_changed
    }

    /// Reload skills if changes are detected
    pub fn reload_if_changed(&self) -> bool {
        if self.has_changed() {
            tracing::info!("Skills directories changed, reloading...");
            self.load_skills();
            true
        } else {
            false
        }
    }

    /// Get the latest modification time of a directory or any of its children
    fn get_dir_mtime(dir: &Path) -> Option<SystemTime> {
        let mut max_mtime = None;

        if let Ok(metadata) = std::fs::metadata(dir)
            && let Ok(mtime) = metadata.modified() {
                max_mtime = Some(mtime);
            }

        if let Ok(entries) = std::fs::read_dir(dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if let Ok(metadata) = std::fs::metadata(&path)
                    && let Ok(mtime) = metadata.modified()
                        && max_mtime.is_none_or(|current| mtime > current) {
                            max_mtime = Some(mtime);
                        }
            }
        }

        max_mtime
    }

    /// Get a copy of loaded skills (checks for changes first)
    pub fn get_loaded_skills(&self) -> Vec<Skill> {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        state.loaded_skills.clone()
    }

    /// Get skills marked as always (checks for changes first)
    pub fn get_always_skills(&self) -> Vec<Skill> {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        state.loaded_skills.iter().filter(|s| s.always).cloned().collect()
    }

    /// Get a specific skill by name (checks for changes first)
    pub fn get_skill(&self, name: &str) -> Option<Skill> {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        state.loaded_skills.iter().find(|s| s.name == name).cloned()
    }

    /// List all skills (name + description) (checks for changes first)
    pub fn list_skills(&self) -> Vec<(String, String)> {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        state.loaded_skills
            .iter()
            .map(|s| (s.name.clone(), s.description.clone()))
            .collect()
    }

    /// Build XML summary of all skills (for progressive disclosure) (checks for changes first)
    pub fn build_skills_summary(&self) -> String {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        
        if state.loaded_skills.is_empty() {
            return String::new();
        }

        let mut lines = vec!["<skills>".to_string()];

        for skill in &state.loaded_skills {
            if skill.always {
                continue;
            }
            lines.push("  <skill>".to_string());
            lines.push(format!("    <name>{}</name>", escape_xml(&skill.name)));
            lines.push(format!(
                "    <description>{}</description>",
                escape_xml(&skill.description)
            ));
            if let Some(path) = &skill.path {
                lines.push(format!("    <path>{}</path>", escape_xml(&path.to_string_lossy())));
            }
            lines.push("  </skill>".to_string());
        }

        lines.push("</skills>".to_string());
        lines.join("\n")
    }

    /// Build prompt for always-injected skills (checks for changes first)
    pub fn build_always_skills_prompt(&self) -> String {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();
        
        let always_skills: Vec<_> = state.loaded_skills.iter().filter(|s| s.always).collect();
        if always_skills.is_empty() {
            return String::new();
        }

        let mut parts = Vec::new();
        for skill in always_skills {
            parts.push(format!("## Skill: {}\n\n{}", skill.name, skill.content));
        }

        parts.join("\n\n---\n\n")
    }

    /// Build full skills prompt: directory conventions, always-skill summary, always-skill content
    pub fn build_skills_prompt(&self) -> String {
        self.reload_if_changed();
        let state = self.state.lock().unwrap();

        if state.loaded_skills.is_empty() {
            return String::new();
        }

        let mut prompt = String::from("## Skills\n\n");

        // Directory conventions
        prompt.push_str("### 目录说明\n\n");
        prompt.push_str("- `~/.agent/skills/` — 外部共享 skill 目录（第三方、系统级 skill）\n");
        prompt.push_str("- `~/.picobot/skills/` — 安装 skill 的默认目录\n");
        prompt.push_str("- `./skills/` — 工作目录下的 skill，picobot 自行创建的 skill 存放于此\n\n");
        prompt.push_str("安装或创建 skill 时请按上述目录规范存放。\n\n");

        // Always skills summary
        let always_skills: Vec<_> = state.loaded_skills.iter().filter(|s| s.always).collect();
        if !always_skills.is_empty() {
            prompt.push_str("### 常用技能\n\n");
            for skill in &always_skills {
                let path_str = skill.path.as_ref()
                    .map(|p| p.to_string_lossy().to_string())
                    .unwrap_or_else(|| "—".to_string());
                prompt.push_str(&format!(
                    "- **{}**: {}  [路径: `{}`]\n",
                    skill.name, skill.description, path_str
                ));
            }
            prompt.push('\n');
        }

        // Usage instructions
        prompt.push_str("### 使用方法\n\n");
        prompt.push_str("- 使用 `get_skill` 工具 action=\"list\" 列出所有可用 skill 及其名称、简介、路径\n");
        prompt.push_str("- 使用 `get_skill` 工具 action=\"get\" 并提供 `skill_name` 获取指定 skill 完整内容\n");

        // Always skills full content
        if !always_skills.is_empty() {
            prompt.push_str("\n---\n\n");
            let mut parts = Vec::new();
            for skill in &always_skills {
                parts.push(format!("## Skill: {}\n\n{}", skill.name, skill.content));
            }
            prompt.push_str(&parts.join("\n\n---\n\n"));
        }

        prompt
    }

    /// Load skills from a specific directory
    fn load_skills_from_dir(&self, dir: &Path) -> Vec<Skill> {
        let mut skills = Vec::new();

        let Ok(entries) = std::fs::read_dir(dir) else {
            tracing::warn!(dir = %dir.display(), "Failed to read skills directory");
            return skills;
        };

        for entry in entries.flatten() {
            let path = entry.path();
            if !path.is_dir() {
                continue;
            }

            let skill_file = path.join("SKILL.md");
            if !skill_file.exists() {
                continue;
            }

            match std::fs::read_to_string(&skill_file) {
                Ok(content) => {
                    match self.parse_skill(&path, &content) {
                        Some(skill) => {
                            tracing::debug!(
                                skill = %skill.name,
                                path = %skill_file.display(),
                                always = skill.always,
                                "Loaded skill"
                            );
                            skills.push(skill);
                        }
                        None => {
                            tracing::warn!(
                                path = %skill_file.display(),
                                "Failed to parse skill"
                            );
                        }
                    }
                }
                Err(e) => {
                    tracing::warn!(
                        path = %skill_file.display(),
                        error = %e,
                        "Failed to read skill file"
                    );
                }
            }
        }

        skills
    }

    /// Parse a skill from markdown content
    fn parse_skill(&self, dir: &Path, content: &str) -> Option<Skill> {
        let (meta, body) = self.parse_skill_markdown(content);

        let name = meta.name.or_else(|| {
            dir.file_name()
                .and_then(|n| n.to_str())
                .map(|s| s.to_string())
        })?;

        let description = meta
            .description
            .unwrap_or_else(|| extract_description(&body));

        Some(Skill {
            name,
            description,
            content: body,
            always: meta.always.unwrap_or(false),
            path: Some(dir.to_path_buf()),
        })
    }

    /// Parse skill markdown, extracting frontmatter and body
    fn parse_skill_markdown(&self, content: &str) -> (SkillMarkdownMeta, String) {
        let normalized = content.replace("\r\n", "\n");

        if let Some(stripped) = normalized.strip_prefix("---\n") {
            if let Some(idx) = stripped.find("\n---\n") {
                let frontmatter = stripped[..idx].to_string();
                let body = stripped[idx + 5..].trim().to_string();
                let meta = self.parse_frontmatter(&frontmatter);
                return (meta, body);
            }
            if let Some(frontmatter) = stripped.strip_suffix("\n---") {
                return (self.parse_frontmatter(frontmatter), String::new());
            }
        }

        (SkillMarkdownMeta::default(), normalized)
    }

    /// Parse simple YAML-like frontmatter
    fn parse_frontmatter(&self, content: &str) -> SkillMarkdownMeta {
        let mut meta = SkillMarkdownMeta::default();

        for line in content.lines() {
            let Some((key, val)) = line.split_once(':') else {
                continue;
            };
            let key = key.trim();
            let val = val.trim().trim_matches('"').trim_matches('\'');

            match key {
                "name" => meta.name = Some(val.to_string()),
                "description" => meta.description = Some(val.to_string()),
                "always" => {
                    meta.always = match val.to_lowercase().as_str() {
                        "true" | "1" | "yes" | "on" => Some(true),
                        "false" | "0" | "no" | "off" => Some(false),
                        _ => None,
                    };
                }
                _ => {}
            }
        }

        meta
    }
}

impl Default for SkillsLoader {
    fn default() -> Self {
        Self::new()
    }
}


/// Extract first non-empty, non-heading line as description
fn extract_description(content: &str) -> String {
    content
        .lines()
        .find(|line| !line.starts_with('#') && !line.trim().is_empty())
        .map(|l| l.trim().to_string())
        .unwrap_or_else(|| "No description".to_string())
}

/// Escape XML special characters
fn escape_xml(s: &str) -> String {
    let mut result = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '&' => result.push_str("&amp;"),
            '<' => result.push_str("&lt;"),
            '>' => result.push_str("&gt;"),
            '"' => result.push_str("&quot;"),
            '\'' => result.push_str("&apos;"),
            _ => result.push(c),
        }
    }
    result
}

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

    #[test]
    fn test_parse_skill_without_frontmatter() {
        let loader = SkillsLoader::new();
        let content = "# My Skill\n\nThis is the content.";
        let (meta, body) = loader.parse_skill_markdown(content);

        assert!(meta.name.is_none());
        assert!(meta.description.is_none());
        assert!(body.contains("My Skill"));
    }

    #[test]
    fn test_parse_skill_with_frontmatter() {
        let loader = SkillsLoader::new();
        let content = r#"---
name: test-skill
description: A test skill
always: true
---
# Test Skill

This is the content.
"#;

        let (meta, body) = loader.parse_skill_markdown(content);

        assert_eq!(meta.name, Some("test-skill".to_string()));
        assert_eq!(meta.description, Some("A test skill".to_string()));
        assert_eq!(meta.always, Some(true));
        assert!(body.contains("Test Skill"));
    }

    #[test]
    fn test_escape_xml() {
        assert_eq!(escape_xml("a & b"), "a &amp; b");
        assert_eq!(escape_xml("<tag>"), "&lt;tag&gt;");
        assert_eq!(escape_xml("\"quote\""), "&quot;quote&quot;");
    }

    #[test]
    fn test_extract_description() {
        assert_eq!(
            extract_description("# Title\n\nFirst line of content."),
            "First line of content."
        );
        assert_eq!(extract_description("# Title"), "No description");
    }
}