tiny-agent 0.1.0

use super::sh_quote;
use crate::{sandbox::Sandbox, tools::ToolOutcome};
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use std::sync::Arc;

#[derive(Clone, Copy, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum WriteFileMode {
    /// 用 content 完整覆盖文件内容
    Write,
    /// 用 new_string 替换 old_string,默认要求 old_string 只出现一次
    Replace,
    /// 在 old_string 后插入 content
    InsertAfter,
    /// 在 old_string 前插入 content
    InsertBefore,
    /// 在文件末尾追加 content
    Append,
}

#[derive(Deserialize, JsonSchema)]
pub struct WriteFileArgs {
    /// 要修改的文件路径
    pub path: String,
    /// 编辑模式: write / replace / insert_after / insert_before / append
    pub mode: Option<WriteFileMode>,
    /// write/insert/append 模式写入的内容
    pub content: Option<String>,
    /// replace/insert 模式的精确匹配文本; replace 默认要求唯一匹配
    pub old_string: Option<String>,
    /// replace 模式的新文本;如果省略则使用 content
    pub new_string: Option<String>,
    /// 是否替换所有匹配;未设置时 replace 默认只允许 1 处匹配
    pub replace_all: Option<bool>,
    /// 期望匹配/替换的次数;设置后实际次数必须完全一致
    pub expected_replacements: Option<usize>,
}

pub async fn write_file(
    sandbox: Arc<dyn Sandbox>,
    args: WriteFileArgs,
) -> Result<ToolOutcome, String> {
    let mode = resolve_mode(&args);
    let next_content = match mode {
        WriteFileMode::Write => required_content(&args)?,
        WriteFileMode::Replace => {
            let current = read_existing_file(sandbox.clone(), &args.path).await?;
            replace_content(&current, &args)?
        }
        WriteFileMode::InsertAfter => {
            let current = read_existing_file(sandbox.clone(), &args.path).await?;
            insert_content(&current, &args, InsertPosition::After)?
        }
        WriteFileMode::InsertBefore => {
            let current = read_existing_file(sandbox.clone(), &args.path).await?;
            insert_content(&current, &args, InsertPosition::Before)?
        }
        WriteFileMode::Append => {
            let current = read_file_or_empty(sandbox.clone(), &args.path).await?;
            format!("{}{}", current, required_content(&args)?)
        }
    };

    write_full_content(sandbox, &args.path, &next_content).await?;
    Ok(json!({
        "result": format!("文件 {} 已更新", args.path),
        "mode": mode_name(mode),
    })
    .into())
}

fn resolve_mode(args: &WriteFileArgs) -> WriteFileMode {
    args.mode.unwrap_or_else(|| {
        if args.old_string.is_some() && (args.new_string.is_some() || args.content.is_some()) {
            WriteFileMode::Replace
        } else {
            WriteFileMode::Write
        }
    })
}

fn required_content(args: &WriteFileArgs) -> Result<String, String> {
    args.content
        .clone()
        .or_else(|| args.new_string.clone())
        .ok_or_else(|| "缺少 content".to_string())
}

fn required_old_string(args: &WriteFileArgs) -> Result<&str, String> {
    let old_string = args
        .old_string
        .as_deref()
        .ok_or_else(|| "缺少 old_string".to_string())?;
    if old_string.is_empty() {
        return Err("old_string 不能为空".to_string());
    }
    Ok(old_string)
}

fn replace_content(current: &str, args: &WriteFileArgs) -> Result<String, String> {
    let old_string = required_old_string(args)?;
    let new_string = args
        .new_string
        .clone()
        .or_else(|| args.content.clone())
        .ok_or_else(|| "缺少 new_string".to_string())?;
    let occurrences = current.matches(old_string).count();
    if occurrences == 0 {
        return Err("old_string 未在文件中找到".to_string());
    }

    let replacements = expected_replacements(args, occurrences);
    if occurrences != replacements {
        return Err(format!(
            "old_string 出现 {} 次,期望 {} 次;未执行替换",
            occurrences, replacements
        ));
    }

    Ok(current.replacen(old_string, &new_string, replacements))
}

#[derive(Clone, Copy)]
enum InsertPosition {
    Before,
    After,
}

fn insert_content(
    current: &str,
    args: &WriteFileArgs,
    position: InsertPosition,
) -> Result<String, String> {
    let old_string = required_old_string(args)?;
    let content = required_content(args)?;
    let occurrences = current.matches(old_string).count();
    if occurrences == 0 {
        return Err("old_string 未在文件中找到".to_string());
    }

    let replacements = expected_replacements(args, occurrences);
    if occurrences != replacements {
        return Err(format!(
            "old_string 出现 {} 次,期望 {} 次;未执行插入",
            occurrences, replacements
        ));
    }

    let replacement = match position {
        InsertPosition::Before => format!("{}{}", content, old_string),
        InsertPosition::After => format!("{}{}", old_string, content),
    };
    Ok(current.replacen(old_string, &replacement, replacements))
}

fn expected_replacements(args: &WriteFileArgs, occurrences: usize) -> usize {
    args.expected_replacements.unwrap_or_else(|| {
        if args.replace_all.unwrap_or(false) {
            occurrences
        } else {
            1
        }
    })
}

async fn read_existing_file(sandbox: Arc<dyn Sandbox>, path: &str) -> Result<String, String> {
    let cmd = format!("cat -- {}", sh_quote(path));
    let out = sandbox.execute(&cmd).await.map_err(|e| e.to_string())?;
    if !out.is_success() {
        return Err(format!("读取文件失败: {}", out.stderr.trim()));
    }
    Ok(out.stdout)
}

async fn read_file_or_empty(sandbox: Arc<dyn Sandbox>, path: &str) -> Result<String, String> {
    let cmd = format!(
        "if [ -e {} ]; then cat -- {}; fi",
        sh_quote(path),
        sh_quote(path)
    );
    let out = sandbox.execute(&cmd).await.map_err(|e| e.to_string())?;
    if !out.is_success() {
        return Err(format!("读取文件失败: {}", out.stderr.trim()));
    }
    Ok(out.stdout)
}

async fn write_full_content(
    sandbox: Arc<dyn Sandbox>,
    path: &str,
    content: &str,
) -> Result<(), String> {
    let cmd = format!(
        "parent=$(dirname -- {}); mkdir -p -- \"$parent\" && printf '%s' {} > {}",
        sh_quote(path),
        sh_quote(content),
        sh_quote(path)
    );
    let out = sandbox.execute(&cmd).await.map_err(|e| e.to_string())?;
    if !out.is_success() {
        return Err(format!("写入文件失败: {}", out.stderr.trim()));
    }
    Ok(())
}

fn mode_name(mode: WriteFileMode) -> &'static str {
    match mode {
        WriteFileMode::Write => "write",
        WriteFileMode::Replace => "replace",
        WriteFileMode::InsertAfter => "insert_after",
        WriteFileMode::InsertBefore => "insert_before",
        WriteFileMode::Append => "append",
    }
}