codecraft/server/src/tools/fs.rs

//! File system tools: read, write, edit, list, glob, grep.

use super::{Tool, ToolCtx, ToolOutput, arg_str, truncate};
use crate::types::ImageData;
use anyhow::{Result, bail};
use async_trait::async_trait;
use base64::Engine;
use serde_json::{Value, json};
use std::path::Path;

const MAX_LINES: usize = 2000;
const MAX_LINE_LEN: usize = 2000;
const MAX_OUTPUT: usize = 60_000;

fn image_type(path: &Path) -> Option<&'static str> {
    match path.extension()?.to_str()?.to_ascii_lowercase().as_str() {
        "png" => Some("image/png"),
        "jpg" | "jpeg" => Some("image/jpeg"),
        "gif" => Some("image/gif"),
        "webp" => Some("image/webp"),
        _ => None,
    }
}

/// Apply only the selected hunks (indices of the 3-line-context groups, as in
/// `unified_diff`) of the change old -> new. Returns (result, total hunks).
pub fn apply_hunks(old: &str, new: &str, selected: &[usize]) -> (String, usize) {
    let diff = similar::TextDiff::from_lines(old, new);
    let groups = diff.grouped_ops(3);
    let total = groups.len();
    let mut out = String::new();
    for op in diff.ops() {
        let in_selected = op.tag() != similar::DiffTag::Equal && groups.iter().enumerate().any(|(i, g)| selected.contains(&i) && g.contains(op));
        let lines: Vec<&str> = if in_selected {
            diff.new_slices()[op.new_range()].to_vec()
        } else {
            diff.old_slices()[op.old_range()].to_vec()
        };
        for l in lines {
            out.push_str(l);
        }
    }
    (out, total)
}

pub fn unified_diff(path: &str, old: &str, new: &str) -> String {
    similar::TextDiff::from_lines(old, new)
        .unified_diff()
        .context_radius(3)
        .header(&format!("a/{path}"), &format!("b/{path}"))
        .to_string()
}

fn rel(ctx: &ToolCtx, p: &Path) -> String {
    p.strip_prefix(&ctx.cwd).unwrap_or(p).display().to_string()
}

// ---------------------------------------------------------------------------

pub struct Read;

#[async_trait]
impl Tool for Read {
    fn name(&self) -> &str {
        "read"
    }
    fn description(&self) -> String {
        "Read a file. Returns numbered lines (cat -n style). Use offset/limit for large files \
         (default: first 2000 lines). Images (png, jpg, gif, webp) are returned as images. \
         Documents are converted to text: PDF (pages=\"3-7\", scanned pages come as images), \
         Excel/ODS spreadsheets (sheet=\"Name\"), Word .docx and PowerPoint .pptx. \
         Always read a file before editing or overwriting it."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "File path, absolute or relative to the working directory"},
                "offset": {"type": "integer", "description": "1-based line to start from"},
                "limit": {"type": "integer", "description": "Number of lines to read"},
                "pages": {"type": "string", "description": "PDF only: page or range, e.g. \"2-5\""},
                "sheet": {"type": "string", "description": "Spreadsheets only: sheet name"}
            },
            "required": ["path"]
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    fn title(&self, input: &Value) -> String {
        format!("Read {}", input["path"].as_str().unwrap_or("?"))
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let path = ctx.resolve_read(arg_str(&input, "path")?)?;
        if path.is_dir() {
            bail!("{} is a directory – use the list tool", path.display());
        }
        if crate::docs::is_document(&path) {
            ctx.app.mark_read(&ctx.session_id, &path);
            let (text, images) = crate::docs::extract(&path, input["pages"].as_str(), input["sheet"].as_str()).await?;
            return Ok(ToolOutput { content: truncate(&text, MAX_OUTPUT * 2), images, ..Default::default() });
        }
        let bytes = tokio::fs::read(&path).await?;
        ctx.app.mark_read(&ctx.session_id, &path);
        if let Some(mt) = image_type(&path) {
            return Ok(ToolOutput {
                content: format!("Image {} ({} bytes)", rel(ctx, &path), bytes.len()),
                images: vec![ImageData {
                    blob: None,
                    size: None,
                    media_type: mt.into(),
                    data: base64::engine::general_purpose::STANDARD.encode(&bytes),
                }],
                ..Default::default()
            });
        }
        if bytes.iter().take(8000).any(|b| *b == 0) {
            return Ok(ToolOutput::text(format!("{} is a binary file ({} bytes)", path.display(), bytes.len())));
        }
        let text = String::from_utf8_lossy(&bytes);
        if text.is_empty() {
            return Ok(ToolOutput::text("(file is empty)"));
        }
        let offset = input["offset"].as_u64().unwrap_or(1).max(1) as usize;
        let limit = input["limit"].as_u64().map(|l| l as usize).unwrap_or(MAX_LINES);
        let total = text.lines().count();
        let mut out = String::new();
        for (i, line) in text.lines().enumerate().skip(offset - 1).take(limit) {
            let line = if line.len() > MAX_LINE_LEN {
                let mut cut = MAX_LINE_LEN;
                while !line.is_char_boundary(cut) {
                    cut -= 1;
                }
                format!("{}…", &line[..cut])
            } else {
                line.to_string()
            };
            out.push_str(&format!("{:>6}\t{}\n", i + 1, line));
        }
        let shown_end = (offset - 1 + limit).min(total);
        if shown_end < total {
            out.push_str(&format!("\n[showing lines {offset}-{shown_end} of {total}; use offset to read more]\n"));
        }
        Ok(ToolOutput::text(truncate(&out, MAX_OUTPUT)))
    }
}

// ---------------------------------------------------------------------------

pub struct Write;

#[async_trait]
impl Tool for Write {
    fn name(&self) -> &str {
        "write"
    }
    fn description(&self) -> String {
        "Create or overwrite a file with the given content. Existing files must be read first. \
         Prefer the edit tool for changes to existing files."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "content": {"type": "string"}
            },
            "required": ["path", "content"]
        })
    }
    fn is_edit(&self) -> bool {
        true
    }
    fn title(&self, input: &Value) -> String {
        format!("Write {}", input["path"].as_str().unwrap_or("?"))
    }
    async fn preview(&self, input: &Value, ctx: &ToolCtx) -> Option<String> {
        let path = ctx.resolve(input["path"].as_str()?).ok()?;
        let old = tokio::fs::read_to_string(&path).await.unwrap_or_default();
        Some(unified_diff(&rel(ctx, &path), &old, input["content"].as_str()?))
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let path = ctx.resolve(arg_str(&input, "path")?)?;
        let content = arg_str(&input, "content")?;
        let old = if path.exists() {
            if !ctx.app.was_read(&ctx.session_id, &path) {
                return Ok(ToolOutput::error("File exists but has not been read in this session. Read it first."));
            }
            tokio::fs::read_to_string(&path).await.unwrap_or_default()
        } else {
            String::new()
        };
        if let Some(dir) = path.parent() {
            tokio::fs::create_dir_all(dir).await?;
        }
        tokio::fs::write(&path, content).await?;
        ctx.app.mark_read(&ctx.session_id, &path);
        let diff = unified_diff(&rel(ctx, &path), &old, content);
        Ok(ToolOutput::text(format!("Wrote {} ({} lines)", path.display(), content.lines().count()))
            .with_meta(json!({"diff": diff, "path": path})))
    }
}

// ---------------------------------------------------------------------------

pub struct Edit;

#[async_trait]
impl Tool for Edit {
    fn name(&self) -> &str {
        "edit"
    }
    fn description(&self) -> String {
        "Replace text in a file. old_string must match the file exactly (including indentation) \
         and be unique unless replace_all is true. The file must have been read first."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "old_string": {"type": "string", "description": "Exact text to replace"},
                "new_string": {"type": "string", "description": "Replacement text"},
                "replace_all": {"type": "boolean", "description": "Replace every occurrence"}
            },
            "required": ["path", "old_string", "new_string"]
        })
    }
    fn is_edit(&self) -> bool {
        true
    }
    fn title(&self, input: &Value) -> String {
        format!("Edit {}", input["path"].as_str().unwrap_or("?"))
    }
    async fn preview(&self, input: &Value, ctx: &ToolCtx) -> Option<String> {
        let path = ctx.resolve(input["path"].as_str()?).ok()?;
        let old = tokio::fs::read_to_string(&path).await.ok()?;
        let new = apply_edit(&old, input).ok()?;
        Some(unified_diff(&rel(ctx, &path), &old, &new))
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let path = ctx.resolve(arg_str(&input, "path")?)?;
        if !ctx.app.was_read(&ctx.session_id, &path) {
            return Ok(ToolOutput::error("File has not been read in this session. Read it first."));
        }
        let old = tokio::fs::read_to_string(&path).await?;
        let new = match apply_edit(&old, &input) {
            Ok(n) => n,
            Err(e) => return Ok(ToolOutput::error(e.to_string())),
        };
        tokio::fs::write(&path, &new).await?;
        let diff = unified_diff(&rel(ctx, &path), &old, &new);
        Ok(ToolOutput::text(format!("Edited {}", path.display()))
            .with_meta(json!({"diff": diff, "path": path})))
    }
}

pub fn apply_edit(old: &str, input: &Value) -> Result<String> {
    let from = arg_str(input, "old_string")?;
    let to = arg_str(input, "new_string")?;
    if from.is_empty() {
        bail!("old_string must not be empty (use write to create files)");
    }
    if from == to {
        bail!("old_string and new_string are identical");
    }
    let count = old.matches(from).count();
    let all = input["replace_all"].as_bool().unwrap_or(false);
    match count {
        0 => bail!("old_string not found in file. Re-read the file and match it exactly."),
        1 => Ok(old.replacen(from, to, 1)),
        n if all => {
            let _ = n;
            Ok(old.replace(from, to))
        }
        n => bail!("old_string occurs {n} times. Add surrounding context to make it unique or set replace_all."),
    }
}

// ---------------------------------------------------------------------------

pub struct List;

#[async_trait]
impl Tool for List {
    fn name(&self) -> &str {
        "list"
    }
    fn description(&self) -> String {
        "List a directory (respects .gitignore). Directories end with '/'.".into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Directory, default: working directory"},
                "depth": {"type": "integer", "description": "Recursion depth (default 1, max 5)"}
            }
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let dir = ctx.resolve_read(input["path"].as_str().unwrap_or("."))?;
        let depth = input["depth"].as_u64().unwrap_or(1).clamp(1, 5) as usize;
        let cwd = ctx.cwd.clone();
        let out = tokio::task::spawn_blocking(move || {
            let mut lines = Vec::new();
            for e in ignore::WalkBuilder::new(&dir).max_depth(Some(depth)).hidden(false).build().flatten() {
                if e.depth() == 0 || e.path().components().any(|c| c.as_os_str() == ".git") {
                    continue;
                }
                let p = e.path().strip_prefix(&dir).unwrap_or(e.path());
                let is_dir = e.file_type().is_some_and(|t| t.is_dir());
                lines.push(format!("{}{}", p.display(), if is_dir { "/" } else { "" }));
                if lines.len() >= 1000 {
                    lines.push("[... truncated at 1000 entries ...]".into());
                    break;
                }
            }
            lines.sort();
            let _ = cwd;
            lines.join("\n")
        })
        .await?;
        Ok(ToolOutput::text(if out.is_empty() { "(empty directory)".into() } else { out }))
    }
}

// ---------------------------------------------------------------------------

pub struct Glob;

#[async_trait]
impl Tool for Glob {
    fn name(&self) -> &str {
        "glob"
    }
    fn description(&self) -> String {
        "Find files by glob pattern (e.g. \"**/*.rs\", \"src/**/test_*.py\"). Respects .gitignore. \
         Results are sorted by modification time, newest first."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "pattern": {"type": "string"},
                "path": {"type": "string", "description": "Directory to search, default: working directory"}
            },
            "required": ["pattern"]
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let pattern = arg_str(&input, "pattern")?.to_string();
        let root = ctx.resolve_read(input["path"].as_str().unwrap_or("."))?;
        let out = tokio::task::spawn_blocking(move || -> Result<String> {
            let matcher = globset::GlobBuilder::new(&pattern).literal_separator(true).build()?.compile_matcher();
            let mut hits = Vec::new();
            for e in ignore::WalkBuilder::new(&root).hidden(false).build().flatten() {
                if !e.file_type().is_some_and(|t| t.is_file()) {
                    continue;
                }
                let relp = e.path().strip_prefix(&root).unwrap_or(e.path());
                if relp.components().any(|c| c.as_os_str() == ".git") {
                    continue;
                }
                if matcher.is_match(relp) {
                    let mtime = e.metadata().ok().and_then(|m| m.modified().ok());
                    hits.push((mtime, e.path().display().to_string()));
                }
            }
            hits.sort_by(|a, b| b.0.cmp(&a.0));
            let n = hits.len();
            let mut s: Vec<String> = hits.into_iter().take(500).map(|h| h.1).collect();
            if n > 500 {
                s.push(format!("[... {} more files ...]", n - 500));
            }
            Ok(if s.is_empty() { "No files found".into() } else { s.join("\n") })
        })
        .await??;
        Ok(ToolOutput::text(out))
    }
}

// ---------------------------------------------------------------------------

pub struct Grep;

#[async_trait]
impl Tool for Grep {
    fn name(&self) -> &str {
        "grep"
    }
    fn description(&self) -> String {
        "Search file contents with a regular expression (Rust regex syntax). Respects .gitignore. \
         output_mode: \"files\" (default, matching file paths), \"content\" (matching lines with \
         line numbers), \"count\"."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "pattern": {"type": "string"},
                "path": {"type": "string", "description": "File or directory, default: working directory"},
                "glob": {"type": "string", "description": "Only search files matching this glob, e.g. \"*.ts\""},
                "output_mode": {"type": "string", "enum": ["files", "content", "count"]},
                "case_insensitive": {"type": "boolean"},
                "context": {"type": "integer", "description": "Lines of context around matches (content mode)"}
            },
            "required": ["pattern"]
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let re = regex::RegexBuilder::new(arg_str(&input, "pattern")?)
            .case_insensitive(input["case_insensitive"].as_bool().unwrap_or(false))
            .build()?;
        let root = ctx.resolve_read(input["path"].as_str().unwrap_or("."))?;
        let glob = input["glob"]
            .as_str()
            .map(|g| globset::Glob::new(g).map(|g| g.compile_matcher()))
            .transpose()?;
        let mode = input["output_mode"].as_str().unwrap_or("files").to_string();
        let context = input["context"].as_u64().unwrap_or(0) as usize;
        let cancel = ctx.cancel.clone();
        let out = tokio::task::spawn_blocking(move || {
            let mut out = String::new();
            let mut files = 0usize;
            for e in ignore::WalkBuilder::new(&root).hidden(false).build().flatten() {
                if cancel.is_cancelled() || out.len() > MAX_OUTPUT {
                    break;
                }
                if !e.file_type().is_some_and(|t| t.is_file()) {
                    continue;
                }
                let p = e.path();
                if p.components().any(|c| c.as_os_str() == ".git") {
                    continue;
                }
                if let Some(g) = &glob {
                    let name = p.file_name().map(Path::new).unwrap_or(p);
                    if !g.is_match(name) && !g.is_match(p.strip_prefix(&root).unwrap_or(p)) {
                        continue;
                    }
                }
                let Ok(bytes) = std::fs::read(p) else { continue };
                if bytes.iter().take(8000).any(|b| *b == 0) {
                    continue;
                }
                let text = String::from_utf8_lossy(&bytes);
                let lines: Vec<&str> = text.lines().collect();
                let hits: Vec<usize> = lines.iter().enumerate().filter(|(_, l)| re.is_match(l)).map(|(i, _)| i).collect();
                if hits.is_empty() {
                    continue;
                }
                files += 1;
                match mode.as_str() {
                    "count" => out.push_str(&format!("{}:{}\n", p.display(), hits.len())),
                    "content" => {
                        let mut last_printed: Option<usize> = None;
                        for &h in &hits {
                            let start = h.saturating_sub(context);
                            let end = (h + context).min(lines.len().saturating_sub(1));
                            for i in start..=end {
                                if last_printed.is_some_and(|l| i <= l) {
                                    continue;
                                }
                                let sep = if i == h { ':' } else { '-' };
                                let line: String = lines[i].chars().take(MAX_LINE_LEN).collect();
                                out.push_str(&format!("{}{sep}{}{sep}{}\n", p.display(), i + 1, line));
                                last_printed = Some(i);
                            }
                        }
                    }
                    _ => out.push_str(&format!("{}\n", p.display())),
                }
            }
            if files == 0 { "No matches found".to_string() } else { out }
        })
        .await?;
        Ok(ToolOutput::text(truncate(&out, MAX_OUTPUT)))
    }
}

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

    #[test]
    fn edit_requires_unique_match() {
        let input = json!({"old_string": "a", "new_string": "b"});
        assert!(apply_edit("a a", &input).is_err());
        assert_eq!(apply_edit("a c", &input).unwrap(), "b c");
        let all = json!({"old_string": "a", "new_string": "b", "replace_all": true});
        assert_eq!(apply_edit("a a", &all).unwrap(), "b b");
        assert!(apply_edit("xyz", &input).is_err());
    }
}

#[cfg(test)]
mod hunk_tests {
    #[test]
    fn partial_apply() {
        let old: String = (1..=30).map(|i| format!("line {i}\n")).collect();
        let new = old.replace("line 2\n", "line TWO\n").replace("line 25\n", "line 25\nextra\n");
        let (only_first, total) = super::apply_hunks(&old, &new, &[0]);
        assert_eq!(total, 2);
        assert!(only_first.contains("line TWO") && !only_first.contains("extra"));
        let (only_second, _) = super::apply_hunks(&old, &new, &[1]);
        assert!(!only_second.contains("line TWO") && only_second.contains("extra"));
        assert_eq!(super::apply_hunks(&old, &new, &[0, 1]).0, new);
        assert_eq!(super::apply_hunks(&old, &new, &[]).0, old);
    }
}