import os
import re

def bake_nvim_safeguards():
    filepath = 'nvim-core/src/lib.rs'
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()

    # 1. We're going to inject an AST-level safeguard inside "nvim_execute_lua"
    # Find the execute_lua block
    target = """                        "nvim_execute_lua" => {
                            if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
                                match execute_nvim_lua(code).await {"""
    
    safeguarded = """                        "nvim_execute_lua" => {
                            if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
                                // BAKE IN: Block interactive prompts that cause server deadlocks
                                let lower_code = code.to_lowercase();
                                if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
                                    send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
                                    return;
                                }
                                // Block the 'c' confirm flag in vim.cmd substitutions
                                if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && code.contains("c") {
                                    // A rough heuristic to block `%s/foo/bar/gc`
                                    if regex::Regex::new(r"%s.*/.*c").unwrap().is_match(code) {
                                        send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/ge' instead.").await;
                                        return;
                                    }
                                }

                                match execute_nvim_lua(code).await {"""
    
    if target in content:
        # Also need to add regex as a dependency or just use basic string matching.
        # Let's just use basic string matching for the %s block to avoid adding the regex crate to nvim-core if not present
        
        safeguarded_no_regex = """                        "nvim_execute_lua" => {
                            if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
                                // BAKE IN: Block interactive prompts that cause server deadlocks
                                let lower_code = code.to_lowercase();
                                if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
                                    send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
                                    return;
                                }
                                // Block the 'c' confirm flag in vim.cmd substitutions
                                if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && (code.contains("gc'") || code.contains('gc"') || code.contains("gc\\n") || code.contains("c'") || code.contains('c"')) {
                                     send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/g' or '/ge' instead.").await;
                                     return;
                                }

                                match execute_nvim_lua(code).await {"""
        
        content = content.replace(target, safeguarded_no_regex)
        print("Baked safeguards into nvim_execute_lua")
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
    else:
        print("Could not find nvim_execute_lua block")

bake_nvim_safeguards()