PicoBot/src/tools/file_edit.rs

use std::path::Path;

use async_trait::async_trait;
use serde_json::json;

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

pub struct FileEditTool {
    allowed_dir: Option<String>,
}

impl FileEditTool {
    pub fn new() -> Self {
        Self { allowed_dir: None }
    }

    pub fn with_allowed_dir(dir: String) -> Self {
        Self {
            allowed_dir: Some(dir),
        }
    }

    fn resolve_path(&self, path: &str) -> Result<std::path::PathBuf, String> {
        let p = Path::new(path);
        let resolved = if p.is_absolute() {
            p.to_path_buf()
        } else {
            std::env::current_dir()
                .map_err(|e| format!("Failed to get current directory: {}", e))?
                .join(p)
        };

        // Check directory restriction
        if let Some(ref allowed) = self.allowed_dir {
            // canonicalize both paths to resolve symlinks and prevent path traversal
            // via symlinks inside allowed_dir pointing outside.
            // For edit tool the target file may not exist yet; fall back to
            // canonicalizing the parent directory.
            let canonical_allowed = std::fs::canonicalize(allowed)
                .map_err(|e| format!("Failed to canonicalize allowed dir '{}': {}", allowed, e))?;
            let canonical_resolved = match std::fs::canonicalize(&resolved) {
                Ok(c) => c,
                Err(_) => {
                    // File doesn't exist yet; canonicalize parent directory
                    let parent = resolved
                        .parent()
                        .ok_or_else(|| format!("Path '{}' has no parent directory", path))?;
                    let canonical_parent = std::fs::canonicalize(parent).map_err(|e| {
                        format!(
                            "Failed to canonicalize parent directory of '{}': {}",
                            path, e
                        )
                    })?;
                    canonical_parent.join(
                        resolved
                            .file_name()
                            .ok_or_else(|| format!("Path '{}' has no file name component", path))?,
                    )
                }
            };
            if !canonical_resolved.starts_with(&canonical_allowed) {
                return Err(format!(
                    "Path '{}' (resolves to '{}') is outside allowed directory '{}'",
                    path,
                    canonical_resolved.display(),
                    canonical_allowed.display()
                ));
            }
        }

        Ok(resolved)
    }

    fn find_match(&self, content: &str, old_text: &str) -> Option<(String, usize)> {
        // Try exact match first
        if content.contains(old_text) {
            let count = content.matches(old_text).count();
            return Some((old_text.to_string(), count));
        }

        // Try line-based matching for minor differences
        let old_lines: Vec<&str> = old_text.lines().collect();
        if old_lines.is_empty() {
            return None;
        }

        let content_lines: Vec<&str> = content.lines().collect();

        for i in 0..content_lines.len().saturating_sub(old_lines.len()) {
            let window = &content_lines[i..i + old_lines.len()];
            let stripped_old: Vec<&str> = old_lines.iter().map(|l| l.trim()).collect();
            let stripped_window: Vec<&str> = window.iter().map(|l| l.trim()).collect();

            if stripped_old == stripped_window {
                let matched_text = window.join("\n");
                return Some((matched_text, 1));
            }
        }

        None
    }
}

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

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

    fn description(&self) -> &str {
        "Edit a file by replacing old_text with new_text. Supports minor whitespace differences. Set replace_all=true to replace every occurrence."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "The file path to edit"
                },
                "old_text": {
                    "type": "string",
                    "description": "The text to find and replace"
                },
                "new_text": {
                    "type": "string",
                    "description": "The text to replace with"
                },
                "replace_all": {
                    "type": "boolean",
                    "description": "Replace all occurrences (default false)",
                    "default": false
                }
            },
            "required": ["path", "old_text", "new_text"]
        })
    }

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

        let old_text = match args.get("old_text").and_then(|v| v.as_str()) {
            Some(t) => t,
            None => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some("Missing required parameter: old_text".to_string()),
                });
            }
        };

        let new_text = match args.get("new_text").and_then(|v| v.as_str()) {
            Some(t) => t,
            None => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some("Missing required parameter: new_text".to_string()),
                });
            }
        };

        let replace_all = extract_bool(&args, "replace_all").unwrap_or(false);

        let resolved = match self.resolve_path(path) {
            Ok(p) => p,
            Err(e) => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(e),
                });
            }
        };

        if !resolved.exists() {
            return Ok(ToolResult {
                success: false,
                output: String::new(),
                error: Some(format!("File not found: {}", path)),
            });
        }

        if !resolved.is_file() {
            return Ok(ToolResult {
                success: false,
                output: String::new(),
                error: Some(format!("Not a file: {}", path)),
            });
        }

        // Read file content
        let content = match std::fs::read_to_string(&resolved) {
            Ok(c) => c,
            Err(e) => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(format!("Failed to read file: {}", e)),
                });
            }
        };

        // Detect line ending style
        let uses_crlf = content.contains("\r\n");
        let norm_content = content.replace("\r\n", "\n");
        let norm_old = old_text.replace("\r\n", "\n");

        // Find match
        let (matched_text, count) = match self.find_match(&norm_content, &norm_old) {
            Some(m) => m,
            None => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(format!(
                        "old_text not found in {}. Verify the file content.",
                        path
                    )),
                });
            }
        };

        // Warn if multiple matches but replace_all is false
        if count > 1 && !replace_all {
            return Ok(ToolResult {
                success: false,
                output: String::new(),
                error: Some(format!(
                    "old_text appears {} times. Provide more context to make it unique, or set replace_all=true.",
                    count
                )),
            });
        }

        // Perform replacement
        let norm_new = new_text.replace("\r\n", "\n");
        let new_content = if replace_all {
            norm_content.replace(&matched_text, &norm_new)
        } else {
            norm_content.replacen(&matched_text, &norm_new, 1)
        };

        // Restore line endings if needed
        let final_content = if uses_crlf {
            new_content.replace("\n", "\r\n")
        } else {
            new_content
        };

        // Write back
        match std::fs::write(&resolved, &final_content) {
            Ok(_) => {
                let replacements = if replace_all { count } else { 1 };
                Ok(ToolResult {
                    success: true,
                    output: format!(
                        "Successfully edited {} ({} replacement{} made)",
                        resolved.display(),
                        replacements,
                        if replacements == 1 { "" } else { "s" }
                    ),
                    error: None,
                })
            }
            Err(e) => Ok(ToolResult {
                success: false,
                output: String::new(),
                error: Some(format!("Failed to write file: {}", e)),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[tokio::test]
    async fn test_edit_simple() {
        let mut file = NamedTempFile::new().unwrap();
        writeln!(file, "Hello World").unwrap();
        writeln!(file, "Test line").unwrap();

        let tool = FileEditTool::new();
        let result = tool
            .execute(json!({
                "path": file.path().to_str().unwrap(),
                "old_text": "Hello World",
                "new_text": "Hello Universe"
            }))
            .await
            .unwrap();

        assert!(result.success);

        let content = std::fs::read_to_string(file.path()).unwrap();
        assert!(content.contains("Hello Universe"));
        assert!(!content.contains("Hello World"));
    }

    #[tokio::test]
    async fn test_edit_replace_all() {
        let mut file = NamedTempFile::new().unwrap();
        writeln!(file, "foo bar").unwrap();
        writeln!(file, "foo baz").unwrap();
        writeln!(file, "foo qux").unwrap();

        let tool = FileEditTool::new();
        let result = tool
            .execute(json!({
                "path": file.path().to_str().unwrap(),
                "old_text": "foo",
                "new_text": "bar",
                "replace_all": true
            }))
            .await
            .unwrap();

        assert!(result.success);

        let content = std::fs::read_to_string(file.path()).unwrap();
        assert!(content.contains("bar bar"));
        assert!(content.contains("bar baz"));
        assert!(content.contains("bar qux"));
    }

    #[tokio::test]
    async fn test_edit_file_not_found() {
        let tool = FileEditTool::new();
        let result = tool
            .execute(json!({
                "path": "/nonexistent/file.txt",
                "old_text": "old",
                "new_text": "new"
            }))
            .await
            .unwrap();

        assert!(!result.success);
        assert!(result.error.unwrap().contains("not found"));
    }

    #[tokio::test]
    async fn test_edit_old_text_not_found() {
        let mut file = NamedTempFile::new().unwrap();
        writeln!(file, "Hello World").unwrap();

        let tool = FileEditTool::new();
        let result = tool
            .execute(json!({
                "path": file.path().to_str().unwrap(),
                "old_text": "NonExistent",
                "new_text": "New"
            }))
            .await
            .unwrap();

        assert!(!result.success);
        assert!(result.error.unwrap().contains("not found"));
    }

    #[tokio::test]
    async fn test_edit_multiline() {
        let mut file = NamedTempFile::new().unwrap();
        writeln!(file, "Line 1").unwrap();
        writeln!(file, "Line 2").unwrap();
        writeln!(file, "Line 3").unwrap();

        let tool = FileEditTool::new();
        let result = tool
            .execute(json!({
                "path": file.path().to_str().unwrap(),
                "old_text": "Line 1\nLine 2",
                "new_text": "New Line 1\nNew Line 2"
            }))
            .await
            .unwrap();

        assert!(result.success);

        let content = std::fs::read_to_string(file.path()).unwrap();
        assert!(content.contains("New Line 1"));
        assert!(content.contains("New Line 2"));
    }
}