cortex-agent 0.2.1

Self-learning AI agent with persistent memory, tools, plugins, and a beautiful terminal UI
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

use serde::Deserialize;
use serde_json::Value;

use crate::tool::ToolSpec;

/// A user-defined plugin loaded from `~/.cortex/plugins/*.yaml`.
#[derive(Debug, Clone, Deserialize)]
pub struct PluginDef {
    pub name: String,
    pub description: String,
    #[serde(default = "default_enabled")]
    pub enabled: bool,
    pub command: String,
    pub parameters: Vec<PluginParam>,
}

fn default_enabled() -> bool { true }

#[derive(Debug, Clone, Deserialize)]
pub struct PluginParam {
    pub name: String,
    #[serde(rename = "type")]
    pub param_type: String,
    pub description: String,
    #[serde(default)]
    pub required: bool,
    #[serde(default)]
    pub default: Option<serde_json::Value>,
}

/// Load all plugins from `~/.cortex/plugins/*.yaml`.
pub fn load_plugins() -> Vec<PluginDef> {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    let plugins_dir = format!("{}/.cortex/plugins", home);
    let dir = Path::new(&plugins_dir);
    if !dir.exists() {
        let _ = std::fs::create_dir_all(dir);
        // Create a sample plugin
        let sample = r#"name: hello-world
description: A sample plugin that echoes a greeting
enabled: true
command: echo "Hello, {name}!"
parameters:
  - name: name
    type: string
    description: Who to greet
    required: true
"#;
        let _ = std::fs::write(dir.join("hello-world.yaml"), sample);
        return vec![];
    }

    let mut plugins = Vec::new();
    let entries = match std::fs::read_dir(dir) {
        Ok(e) => e,
        Err(_) => return plugins,
    };

    for entry in entries.flatten() {
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) != Some("yaml") {
            continue;
        }
        let content = match std::fs::read_to_string(&path) {
            Ok(c) => c,
            Err(_) => continue,
        };
        match serde_yaml::from_str::<PluginDef>(&content) {
            Ok(plugin) => {
                if plugin.enabled {
                    plugins.push(plugin);
                }
            }
            Err(e) => {
                eprintln!("  [plugin] Failed to parse {}: {}", path.display(), e);
            }
        }
    }
    plugins
}

/// Convert a PluginDef into a ToolSpec that runs a shell command.
pub fn plugin_to_tool(plugin: &PluginDef) -> ToolSpec {
    let mut properties = serde_json::Map::new();
    let mut required = Vec::new();

    for param in &plugin.parameters {
        let schema = serde_json::json!({
            "type": param.param_type,
            "description": param.description,
        });
        let mut schema = if let Some(default) = &param.default {
            let mut m = schema.as_object().unwrap().clone();
            m.insert("default".into(), default.clone());
            Value::Object(m)
        } else {
            schema
        };
        properties.insert(param.name.clone(), schema);
        if param.required {
            required.push(param.name.clone());
        }
    }

    let parameters = serde_json::json!({
        "type": "object",
        "properties": properties,
        "required": required,
    });

    let command_template = plugin.command.clone();
    let param_names: Vec<String> = plugin.parameters.iter().map(|p| p.name.clone()).collect();

    ToolSpec::new(
        plugin.name.clone(),
        plugin.description.clone(),
        parameters,
        Arc::new(move |args| {
            let mut cmd_str = command_template.clone();
            for pname in &param_names {
                let val = args.get(pname)
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string())
                    .or_else(|| args.get(pname).map(|v| v.to_string()))
                    .unwrap_or_default();
                cmd_str = cmd_str.replace(&format!("{{{}}}", pname), &val);
            }
            let output = std::process::Command::new("bash")
                .arg("-c")
                .arg(&cmd_str)
                .output()
                .map_err(|e| format!("Plugin execution error: {}", e))?;
            let stdout = String::from_utf8_lossy(&output.stdout).to_string();
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            let mut result = String::new();
            if !stdout.is_empty() { result.push_str(&stdout); }
            if !stderr.is_empty() { result.push_str(&format!("\n[stderr]\n{}", stderr)); }
            if result.is_empty() { result = format!("(exit code: {})", output.status.code().unwrap_or(-1)); }
            Ok(result.trim().to_string())
        }),
    )
}

/// Get plugin state (enabled/disabled) from the plugins directory.
pub fn list_plugin_states() -> Vec<(String, bool)> {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    let plugins_dir = format!("{}/.cortex/plugins", home);
    let dir = Path::new(&plugins_dir);
    if !dir.exists() {
        return vec![];
    }
    let mut states = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("yaml") {
                continue;
            }
            if let Ok(content) = std::fs::read_to_string(&path) {
                if let Ok(plugin) = serde_yaml::from_str::<PluginDef>(&content) {
                    states.push((plugin.name, plugin.enabled));
                }
            }
        }
    }
    states
}

/// Enable or disable a plugin by toggling its `enabled` field.
pub fn set_plugin_enabled(name: &str, enabled: bool) -> Result<(), String> {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    let plugins_dir = format!("{}/.cortex/plugins", home);
    let dir = Path::new(&plugins_dir);
    if !dir.exists() {
        return Err("Plugins directory not found".into());
    }
    for entry in std::fs::read_dir(dir).map_err(|e| e.to_string())? {
        let entry = entry.map_err(|e| e.to_string())?;
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) != Some("yaml") {
            continue;
        }
        let content = std::fs::read_to_string(&path).map_err(|e| e.to_string())?;
        if let Ok(plugin) = serde_yaml::from_str::<PluginDef>(&content) {
            if plugin.name == name {
                let mut lines: Vec<String> = content.lines().map(|l| l.to_string()).collect();
                let mut found = false;
                for line in &mut lines {
                    if line.trim().starts_with("enabled:") {
                        *line = format!("enabled: {}", enabled);
                        found = true;
                        break;
                    }
                }
                if !found {
                    lines.push(format!("enabled: {}", enabled));
                }
                let new_content = lines.join("\n");
                std::fs::write(&path, &new_content).map_err(|e| e.to_string())?;
                return Ok(());
            }
        }
    }
    Err(format!("Plugin '{}' not found", name))
}

/// Download a plugin from a URL and install it to ~/.cortex/plugins/.
pub fn install_plugin_from_url(url: &str) -> Result<String, String> {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    let plugins_dir = format!("{}/.cortex/plugins", home);
    let dir = std::path::Path::new(&plugins_dir);
    let _ = std::fs::create_dir_all(dir);

    let yaml_url = if url.contains("github.com") && !url.contains("raw.githubusercontent.com") {
        let repo_part = url
            .trim_start_matches("https://github.com/")
            .trim_start_matches("http://github.com/")
            .trim_end_matches('/');
        let parts: Vec<&str> = repo_part.splitn(2, '/').collect();
        if parts.len() < 2 {
            return Err("Invalid GitHub URL. Use: https://github.com/user/repo or a raw URL".into());
        }
        let (user, repo_and_path) = (parts[0], parts[1]);
        let path = repo_and_path.strip_prefix("blob/").unwrap_or(repo_and_path);
        if path.contains('/') {
            format!("https://raw.githubusercontent.com/{}/{}", user, path)
        } else {
            let repo = repo_and_path.trim_end_matches(".git");
            let candidates = ["plugin.yaml", "plugin.yml", "cortex-plugin.yaml"];
            let mut found = None;
            for candidate in &candidates {
                let test_url = format!("https://raw.githubusercontent.com/{}/{}/main/{}", user, repo, candidate);
                let client = reqwest::blocking::Client::builder()
                    .timeout(std::time::Duration::from_secs(5))
                    .build().map_err(|e| format!("HTTP client error: {}", e))?;
                if let Ok(r) = client.head(&test_url).send() {
                    if r.status().is_success() {
                        found = Some(test_url);
                        break;
                    }
                }
            }
            found.ok_or_else(|| format!("No plugin file found in {}. Use full raw URL.", url))?
        }
    } else {
        url.to_string()
    };

    let client = reqwest::blocking::Client::builder()
        .timeout(std::time::Duration::from_secs(15))
        .build().map_err(|e| format!("HTTP client error: {}", e))?;
    let resp = client.get(&yaml_url)
        .header("User-Agent", "Cortex/1.0")
        .send().map_err(|e| format!("Failed to download: {}", e))?;
    if !resp.status().is_success() {
        return Err(format!("HTTP {}: {}", resp.status(), yaml_url));
    }
    let content = resp.text().map_err(|e| format!("Failed to read response: {}", e))?;

    let plugin: PluginDef = serde_yaml::from_str(&content)
        .map_err(|e| format!("Invalid plugin YAML: {}", e))?;
    if plugin.name.is_empty() {
        return Err("Plugin must have a 'name' field".into());
    }

    let filename = format!("{}.yaml", plugin.name);
    let dest = dir.join(&filename);
    std::fs::write(&dest, &content).map_err(|e| format!("Failed to save: {}", e))?;
    Ok(format!("Installed plugin '{}' from {}", plugin.name, yaml_url))
}