magi-code 0.96.1

use crate::{tools::ToolResult, tools::contract::tool_name};
use serde_json::Value;
use std::collections::HashSet;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ContextInjectionDefinition {
    pub(crate) name: &'static str,
    pub(crate) tag: &'static str,
    pub(crate) command: &'static str,
    pub(crate) description: &'static str,
}

const DEFINITIONS: [ContextInjectionDefinition; 3] = [
    ContextInjectionDefinition {
        name: "tree",
        tag: "#tree",
        command: "tree",
        description: "current directory tree",
    },
    ContextInjectionDefinition {
        name: "git-status",
        tag: "#git-status",
        command: "git status",
        description: "current Git status",
    },
    ContextInjectionDefinition {
        name: "diff-changes",
        tag: "#diff-changes",
        command: "",
        description: "unresolved Diff comments",
    },
];

pub(crate) fn definitions() -> &'static [ContextInjectionDefinition] {
    &DEFINITIONS
}

pub(crate) fn injection_tags_in_prompt(prompt: &str) -> Vec<ContextInjectionDefinition> {
    let mut seen = HashSet::new();
    let mut matches = Vec::new();
    for (index, _) in prompt.match_indices('#') {
        for definition in definitions() {
            if standalone_tag_at(prompt, index, definition.tag) && seen.insert(definition.name) {
                matches.push(*definition);
            }
        }
    }
    matches
}

pub(crate) fn expand_prompt_with_context_injections(
    prompt: &str,
    mut run: impl FnMut(&ContextInjectionDefinition) -> ToolResult,
) -> Option<String> {
    let definitions = injection_tags_in_prompt(prompt);
    if definitions.is_empty() {
        return None;
    }
    let results = definitions
        .iter()
        .map(|definition| (*definition, run(definition)))
        .collect::<Vec<_>>();
    Some(render_effective_prompt(prompt, &results))
}

/// Runs on the agent worker, never while the editor is matching tags.
pub(super) fn expand_user_prompt(
    prompt: &str,
    cwd: &std::path::Path,
    cancellation: &crate::cancellation::AgentCancellation,
    run: impl FnMut(&ContextInjectionDefinition) -> ToolResult,
) -> Option<String> {
    expand_user_prompt_with_diff_context(prompt, run, || {
        crate::diff_review::comments_context(cwd, cancellation)
    })
}

fn expand_user_prompt_with_diff_context(
    prompt: &str,
    mut run: impl FnMut(&ContextInjectionDefinition) -> ToolResult,
    mut diff_context: impl FnMut() -> Result<String, String>,
) -> Option<String> {
    expand_prompt_with_context_injections(prompt, |definition| {
        if definition.name != "diff-changes" {
            return run(definition);
        }
        let context = diff_context();
        ToolResult {
            tool_name: "diff_review".to_string(),
            success: context.is_ok(),
            content: context.unwrap_or_else(|error| format!("Diff comments unavailable: {error}")),
            metadata: Value::Null,
            display: Default::default(),
        }
    })
}

fn standalone_tag_at(prompt: &str, start: usize, tag: &str) -> bool {
    let end = start + tag.len();
    prompt
        .as_bytes()
        .get(start..end)
        .is_some_and(|candidate| candidate == tag.as_bytes())
        && prompt[..start]
            .bytes()
            .next_back()
            .is_none_or(|byte| !is_identifier_byte(byte))
        && prompt[end..]
            .bytes()
            .next()
            .is_none_or(|byte| !is_identifier_byte(byte))
}

fn is_identifier_byte(byte: u8) -> bool {
    byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')
}

fn render_effective_prompt(
    prompt: &str,
    results: &[(ContextInjectionDefinition, ToolResult)],
) -> String {
    let mut rendered = prompt.to_string();
    rendered.push_str("\n\n<context_injections>");
    for (definition, result) in results {
        rendered.push('\n');
        rendered.push_str(&format!(
            "<context_injection name=\"{}\" tag=\"{}\" command=\"{}\" success=\"{}\">\n",
            definition.name, definition.tag, definition.command, result.success
        ));
        if definition.name == "diff-changes" {
            rendered.push_str(&result.content);
            rendered.push('\n');
        } else if result.tool_name != tool_name::BASH {
            rendered.push_str("tool_error:\n");
            rendered.push_str(&result.content);
            rendered.push('\n');
        } else {
            rendered.push_str("stdout:\n");
            rendered.push_str(metadata_string(&result.metadata, "stdout").unwrap_or_default());
            rendered.push_str("\nstderr:\n");
            rendered.push_str(metadata_string(&result.metadata, "stderr").unwrap_or_default());
            rendered.push_str("\nstatus:\n");
            rendered.push_str(&format!("success: {}\n", result.success));
            rendered.push_str(&format!(
                "exit_code: {}\n",
                metadata_exit_code(&result.metadata).unwrap_or_else(|| "unknown".to_string())
            ));
            rendered.push_str(&format!(
                "timed_out: {}\n",
                metadata_bool(&result.metadata, "timed_out")
                    .map(|value| value.to_string())
                    .unwrap_or_else(|| "unknown".to_string())
            ));
            rendered.push_str(&format!(
                "stdout_truncated: {}\n",
                metadata_bool(&result.metadata, "stdout_truncated")
                    .map(|value| value.to_string())
                    .unwrap_or_else(|| "unknown".to_string())
            ));
            rendered.push_str(&format!(
                "stderr_truncated: {}\n",
                metadata_bool(&result.metadata, "stderr_truncated")
                    .map(|value| value.to_string())
                    .unwrap_or_else(|| "unknown".to_string())
            ));
            if !result.success && result.content != "stdout:\n\nstderr:\n" {
                rendered.push_str("content:\n");
                rendered.push_str(&result.content);
                rendered.push('\n');
            }
        }
        rendered.push_str("</context_injection>");
    }
    rendered.push_str("\n</context_injections>");
    rendered
}

fn metadata_string<'a>(metadata: &'a Value, key: &str) -> Option<&'a str> {
    metadata.get(key)?.as_str()
}

fn metadata_bool(metadata: &Value, key: &str) -> Option<bool> {
    metadata.get(key)?.as_bool()
}

fn metadata_exit_code(metadata: &Value) -> Option<String> {
    let value = metadata.get("exit_code")?;
    if value.is_null() {
        return Some("none".to_string());
    }
    value.as_i64().map(|code| code.to_string())
}