PicoBot/src/tools/file_search.rs

use std::path::Path;
use std::process::Stdio;

use async_trait::async_trait;
use serde_json::json;
use tokio::process::Command;
use tokio::time::timeout;

use crate::tools::traits::{Tool, ToolResult};

const MAX_RESULTS: usize = 200;
const MAX_OUTPUT_CHARS: usize = 50_000;
const TIMEOUT_SECS: u64 = 60;

pub struct FileSearchTool;

impl FileSearchTool {
    pub fn new() -> Self {
        Self
    }

    fn resolve_dir(&self, dir: Option<&str>) -> String {
        match dir {
            Some(d) if !d.is_empty() => d.to_string(),
            _ => ".".to_string(),
        }
    }

    fn truncate_output(&self, lines: &[String]) -> String {
        let mut output = String::new();
        for (i, line) in lines.iter().enumerate() {
            if output.len() + line.len() + 1 > MAX_OUTPUT_CHARS {
                let omitted = lines.len() - i;
                output.push_str(&format!("\n... ({} files omitted) ...", omitted));
                break;
            }
            if !output.is_empty() {
                output.push('\n');
            }
            output.push_str(line);
        }
        output
    }
}

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

#[async_trait]
impl Tool for FileSearchTool {
    fn name(&self) -> &str {
        "file_search"
    }

    fn description(&self) -> &str {
        "Search for files by glob pattern (e.g. '*.rs', 'test_*.rs'). Internally uses fd (fast find) for efficient searching — if fd is not available, falls back to the find command, then pure Rust glob matching."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "pattern": {
                    "type": "string",
                    "description": "File glob pattern to search for (e.g. *.rs, test_*.rs, src/**/*.py)"
                },
                "dir": {
                    "type": "string",
                    "description": "Directory to search in (default: current working directory)"
                },
                "case_sensitive": {
                    "type": "boolean",
                    "description": "Whether to match case-sensitively (default: true)"
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of results to return (default: 200)"
                }
            },
            "required": ["pattern"]
        })
    }

    fn read_only(&self) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let pattern = match args.get("pattern").and_then(|v| v.as_str()) {
            Some(p) if !p.is_empty() => p,
            _ => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some("Missing required parameter: pattern".to_string()),
                });
            }
        };

        let dir = self.resolve_dir(args.get("dir").and_then(|v| v.as_str()));
        let case_sensitive = args
            .get("case_sensitive")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);
        let max_results = args
            .get("max_results")
            .and_then(|v| v.as_u64())
            .unwrap_or(MAX_RESULTS as u64) as usize;

        let result = self
            .run_search(pattern, &dir, case_sensitive, max_results)
            .await;

        match result {
            Ok(lines) => {
                let count = lines.len();
                let mut output = self.truncate_output(&lines);
                output.push_str(&format!("\n\n---\n共 {} 个文件", count));
                Ok(ToolResult {
                    success: true,
                    output,
                    error: None,
                })
            }
            Err(e) => Ok(ToolResult {
                success: false,
                output: String::new(),
                error: Some(e.to_string()),
            }),
        }
    }
}

impl FileSearchTool {
    async fn run_search(
        &self,
        pattern: &str,
        dir: &str,
        case_sensitive: bool,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        let fd_cmd = if which::which("fd").is_ok() {
            "fd"
        } else if which::which("fdfind").is_ok() {
            "fdfind"
        } else {
            ""
        };

        if !fd_cmd.is_empty() {
            match self
                .search_with_fd(pattern, dir, case_sensitive, max_results, fd_cmd)
                .await
            {
                Ok(lines) if !lines.is_empty() => return Ok(lines),
                Ok(_) => {}
                Err(e) => tracing::warn!("{} failed: {}, falling back", fd_cmd, e),
            }
        }

        if which::which("find").is_ok() {
            match self.search_with_find(pattern, dir, max_results).await {
                Ok(lines) if !lines.is_empty() => return Ok(lines),
                Ok(_) => {}
                Err(e) => tracing::warn!("find failed: {}, falling back", e),
            }
        }

        tracing::warn!("No fd/find available, using built-in file search (slower)");
        self.search_with_rust(pattern, dir, case_sensitive, max_results)
            .await
    }

    async fn search_with_fd(
        &self,
        pattern: &str,
        dir: &str,
        case_sensitive: bool,
        max_results: usize,
        fd_cmd: &str,
    ) -> anyhow::Result<Vec<String>> {
        let mut cmd = Command::new(fd_cmd);
        cmd.arg("--search-path")
            .arg(dir)
            .arg("--glob")
            .arg(pattern)
            .arg("--color")
            .arg("never")
            .arg("--max-results")
            .arg(max_results.to_string())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        if !case_sensitive {
            cmd.arg("--ignore-case");
        }

        let output = timeout(std::time::Duration::from_secs(TIMEOUT_SECS), cmd.output())
            .await
            .map_err(|_| anyhow::anyhow!("fd timed out after {}s", TIMEOUT_SECS))??;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(anyhow::anyhow!("fd error: {}", stderr.trim()));
        }

        let text = String::from_utf8_lossy(&output.stdout);
        let lines: Vec<String> = text
            .lines()
            .filter(|l| !l.is_empty())
            .map(|l| l.to_string())
            .collect();
        Ok(lines)
    }

    async fn search_with_find(
        &self,
        pattern: &str,
        dir: &str,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        let mut cmd = Command::new("find");
        cmd.arg(dir)
            .arg("-name")
            .arg(pattern)
            .arg("-not")
            .arg("-path")
            .arg("*/.*")
            .stdout(Stdio::piped())
            .stderr(Stdio::null());

        let output = timeout(std::time::Duration::from_secs(TIMEOUT_SECS), cmd.output())
            .await
            .map_err(|_| anyhow::anyhow!("find timed out after {}s", TIMEOUT_SECS))??;

        let text = String::from_utf8_lossy(&output.stdout);
        let mut lines: Vec<String> = text
            .lines()
            .filter(|l| !l.is_empty())
            .map(|l| {
                let p = Path::new(l);
                p.to_string_lossy().to_string()
            })
            .collect();
        if lines.len() > max_results {
            lines.truncate(max_results);
        }
        Ok(lines)
    }

    async fn search_with_rust(
        &self,
        pattern: &str,
        dir: &str,
        case_sensitive: bool,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        let regex_str = glob_to_regex(pattern);
        let re = if case_sensitive {
            regex::Regex::new(&regex_str)
        } else {
            regex::RegexBuilder::new(&regex_str)
                .case_insensitive(true)
                .build()
        }
        .map_err(|e| anyhow::anyhow!("Invalid glob pattern '{}': {}", pattern, e))?;

        let mut results = Vec::new();
        walk_dir(
            Path::new(dir),
            Path::new(dir),
            &re,
            &mut results,
            max_results,
        )?;
        Ok(results)
    }
}

fn glob_to_regex(glob: &str) -> String {
    let mut regex = String::from("^");
    let chars: Vec<char> = glob.chars().collect();
    let mut i = 0;
    while i < chars.len() {
        match chars[i] {
            '*' => {
                if i + 1 < chars.len() && chars[i + 1] == '*' {
                    regex.push_str(".*");
                    i += 1;
                } else {
                    regex.push_str("[^/]*");
                }
            }
            '?' => regex.push_str("[^/]"),
            '.' | '+' | '(' | ')' | '[' | ']' | '{' | '}' | '^' | '$' | '|' | '\\' => {
                regex.push('\\');
                regex.push(chars[i]);
            }
            c => regex.push(c),
        }
        i += 1;
    }
    regex.push('$');
    regex
}

fn walk_dir(
    base: &Path,
    current: &Path,
    re: &regex::Regex,
    results: &mut Vec<String>,
    max: usize,
) -> anyhow::Result<()> {
    if results.len() >= max {
        return Ok(());
    }

    let entries = match std::fs::read_dir(current) {
        Ok(e) => e,
        Err(_) => return Ok(()),
    };

    for entry in entries.flatten() {
        let path = entry.path();
        let rel = match path.strip_prefix(base) {
            Ok(r) => r,
            Err(_) => continue,
        };

        if path.is_dir() {
            if let Some(name) = rel.file_name().and_then(|n| n.to_str())
                && name.starts_with('.')
                && name.len() > 1
            {
                continue;
            }
            walk_dir(base, &path, re, results, max)?;
        } else if path.is_file() {
            if let Some(name) = rel.file_name().and_then(|n| n.to_str())
                && re.is_match(name)
            {
                results.push(rel.to_string_lossy().to_string());
            }
            if results.len() >= max {
                return Ok(());
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    #[cfg(unix)]
    #[tokio::test]
    async fn fd_search_path_does_not_use_incompatible_strip_cwd_prefix() {
        use std::os::unix::fs::PermissionsExt;

        let dir = TempDir::new().unwrap();
        let fake_fd = dir.path().join("fd");
        fs::write(
            &fake_fd,
            r#"#!/bin/sh
for arg in "$@"; do
    if [ "$arg" = "--strip-cwd-prefix" ]; then
        echo "--strip-cwd-prefix conflicts with --search-path" >&2
        exit 2
    fi
done
printf '%s\n' 'example.rs'
"#,
        )
        .unwrap();
        let mut permissions = fs::metadata(&fake_fd).unwrap().permissions();
        permissions.set_mode(0o755);
        fs::set_permissions(&fake_fd, permissions).unwrap();

        let tool = FileSearchTool::new();
        let results = tool
            .search_with_fd(
                "*.rs",
                dir.path().to_str().unwrap(),
                true,
                10,
                fake_fd.to_str().unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(results, vec!["example.rs"]);
    }

    #[tokio::test]
    async fn test_file_search_rust_fallback() {
        let dir = TempDir::new().unwrap();
        fs::write(dir.path().join("main.rs"), "fn main() {}").unwrap();
        fs::write(dir.path().join("lib.rs"), "pub fn foo() {}").unwrap();
        fs::write(dir.path().join("test.rs"), "#[test] fn t() {}").unwrap();
        fs::write(dir.path().join("README.md"), "# Readme").unwrap();
        fs::create_dir(dir.path().join("src")).unwrap();
        fs::write(dir.path().join("src/nested.rs"), "fn nested() {}").unwrap();

        let tool = FileSearchTool::new();
        let result = tool
            .execute(json!({
                "pattern": "*.rs",
                "dir": dir.path().to_str().unwrap()
            }))
            .await
            .unwrap();

        assert!(result.success);
        assert!(result.output.contains("main.rs"));
        assert!(result.output.contains("lib.rs"));
        assert!(result.output.contains("test.rs"));
        assert!(result.output.contains("nested.rs"));
        assert!(!result.output.contains("README.md"));
        assert!(result.output.contains("共 4 个文件"));
    }

    #[tokio::test]
    async fn test_file_search_max_results() {
        let dir = TempDir::new().unwrap();
        for i in 0..5 {
            fs::write(dir.path().join(format!("file_{}.rs", i)), "").unwrap();
        }

        let tool = FileSearchTool::new();
        let result = tool
            .execute(json!({
                "pattern": "*.rs",
                "dir": dir.path().to_str().unwrap(),
                "max_results": 3
            }))
            .await
            .unwrap();

        assert!(result.success);
        assert!(result.output.contains("共 3 个文件"));
    }
}