PicoBot/src/tools/content_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 = 100;
const MAX_OUTPUT_CHARS: usize = 50_000;
const TIMEOUT_SECS: u64 = 60;

pub struct ContentSearchTool;

impl ContentSearchTool {
    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... ({} matches omitted) ...",
                    omitted
                ));
                break;
            }
            if !output.is_empty() {
                output.push('\n');
            }
            output.push_str(line);
        }
        output
    }
}

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

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

    fn description(&self) -> &str {
        "Search file contents by regex or text pattern. Internally uses ripgrep (rg) for fast searching — if rg is not available, falls back to grep, then pure Rust. Supports context lines, file filtering, and case-sensitivity options."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "pattern": {
                    "type": "string",
                    "description": "Regex or text pattern to search for in file contents"
                },
                "dir": {
                    "type": "string",
                    "description": "Directory to search in (default: current working directory)"
                },
                "file_pattern": {
                    "type": "string",
                    "description": "Optional glob to restrict which files to search (e.g. '*.rs', '*.{rs,toml}')"
                },
                "case_sensitive": {
                    "type": "boolean",
                    "description": "Whether to match case-sensitively (default: false)"
                },
                "context_lines": {
                    "type": "integer",
                    "description": "Number of context lines to show before and after each match (default: 0)"
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of matching lines to return (default: 100)"
                }
            },
            "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 file_pattern = args.get("file_pattern").and_then(|v| v.as_str());
        let case_sensitive = args.get("case_sensitive").and_then(|v| v.as_bool()).unwrap_or(false);
        let context_lines = args.get("context_lines").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
        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, file_pattern, case_sensitive, context_lines, 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 ContentSearchTool {
    async fn run_search(
        &self,
        pattern: &str,
        dir: &str,
        file_pattern: Option<&str>,
        case_sensitive: bool,
        context_lines: usize,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        if which::which("rg").is_ok() {
            match self.search_with_rg(pattern, dir, file_pattern, case_sensitive, context_lines, max_results).await {
                Ok(lines) => return Ok(lines),
                Err(e) => tracing::warn!("rg failed: {}, falling back", e),
            }
        }

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

        tracing::warn!("No rg/grep available, using built-in content search (much slower). Install ripgrep for better performance.");
        self.search_with_rust(pattern, dir, file_pattern, case_sensitive, context_lines, max_results).await
    }

    async fn search_with_rg(
        &self,
        pattern: &str,
        dir: &str,
        file_pattern: Option<&str>,
        case_sensitive: bool,
        context_lines: usize,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        let mut cmd = Command::new("rg");
        cmd.arg("-n")
            .arg("--no-heading")
            .arg("--color").arg("never")
            .arg("--max-count").arg(max_results.to_string())
            .arg(pattern)
            .arg(dir)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        if !case_sensitive {
            cmd.arg("-i");
        }
        if context_lines > 0 {
            cmd.arg("-C").arg(context_lines.to_string());
        }
        if let Some(fp) = file_pattern {
            cmd.arg("--glob").arg(fp);
        }

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

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

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

    async fn search_with_grep(
        &self,
        pattern: &str,
        dir: &str,
        file_pattern: Option<&str>,
        case_sensitive: bool,
        context_lines: usize,
        max_results: usize,
    ) -> anyhow::Result<Vec<String>> {
        let mut cmd = Command::new("grep");
        cmd.arg("-rn")
            .arg("-E")
            .arg("--color=never")
            .arg("--binary-files=without-match")
            .arg(pattern)
            .arg(dir)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        if !case_sensitive {
            cmd.arg("-i");
        }
        if context_lines > 0 {
            cmd.arg("-C").arg(context_lines.to_string());
        }
        if let Some(fp) = file_pattern {
            cmd.arg("--include").arg(fp);
        }

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

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

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

        let file_re = file_pattern.map(|fp| {
            let re_str = glob_to_regex(fp);
            if case_sensitive {
                regex::Regex::new(&re_str)
            } else {
                regex::RegexBuilder::new(&re_str).case_insensitive(true).build()
            }
        });

        let file_re = match file_re {
            Some(Ok(r)) => Some(r),
            Some(Err(e)) => return Err(anyhow::anyhow!("Invalid file pattern: {}", e)),
            None => None,
        };

        let mut results = Vec::new();
        grep_dir(Path::new(dir), Path::new(dir), &re, file_re.as_ref(), &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 grep_dir(
    base: &Path,
    current: &Path,
    re: &regex::Regex,
    file_re: Option<&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;
                }
            grep_dir(base, &path, re, file_re, results, max)?;
        } else if path.is_file() {
            if let Some(file_re) = file_re
                && let Some(name) = rel.file_name().and_then(|n| n.to_str())
                    && !file_re.is_match(name) {
                        continue;
                    }

            if let Ok(content) = std::fs::read_to_string(&path) {
                for (line_num, line) in content.lines().enumerate() {
                    if re.is_match(line) {
                        results.push(format!(
                            "{}:{}:{}",
                            rel.to_string_lossy(),
                            line_num + 1,
                            line
                        ));
                        if results.len() >= max {
                            return Ok(());
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

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

    #[tokio::test]
    async fn test_content_search_rust_fallback() {
        let dir = TempDir::new().unwrap();
        fs::write(dir.path().join("main.rs"), "fn main() {\n    let x = 42;\n    println!(\"hello\");\n}").unwrap();
        fs::write(dir.path().join("lib.rs"), "pub fn foo() -> u32 {\n    let y = 42;\n    y\n}").unwrap();
        fs::write(dir.path().join("README.md"), "# Project\nHello world").unwrap();

        let tool = ContentSearchTool::new();
        let result = tool
            .execute(json!({
                "pattern": "let.*=.*42",
                "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("README.md"));
        assert!(result.output.contains("共 2 条匹配"));
    }

    #[tokio::test]
    async fn test_content_search_file_filter() {
        let dir = TempDir::new().unwrap();
        fs::write(dir.path().join("main.rs"), "fn main() {}").unwrap();
        fs::write(dir.path().join("config.toml"), "name = \"test\"").unwrap();

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

        assert!(result.success);
        assert!(result.output.contains("config.toml"));
        assert!(!result.output.contains("main.rs"));
    }

    #[tokio::test]
    async fn test_content_search_max_results() {
        let dir = TempDir::new().unwrap();
        let mut content = String::new();
        for i in 0..10 {
            content.push_str(&format!("match line {}\n", i));
        }
        fs::write(dir.path().join("data.txt"), &content).unwrap();

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

        assert!(result.success);
        assert!(result.output.contains("共 3 条匹配"));
    }
}