vibe-coded-script.ts

import { readdirSync } from "node:fs";
import { readFile, writeFile } from "node:fs/promises";
import { join, basename } from "node:path";
import { execSync } from "node:child_process";
import { GoogleGenerativeAI } from "@google/generative-ai";
import "dotenv";

const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
const MODEL = "gemini-2.5-pro-preview-06-05	";

if (!GEMINI_API_KEY) {
  console.error("Error: GEMINI_API_KEY environment variable is required");
  console.error("Please add GEMINI_API_KEY=your_api_key to your .env file");
  process.exit(1);
}

interface ResourceConfig {
  name: string;
  paths: string[];
  extensions: string[];
}

const RESOURCES: ResourceConfig[] = [
  {
    name: "evmone",
    paths: ["evmone/**/*.cpp", "evmone/**/*.hpp"],
    extensions: [".cpp", ".hpp"],
  },
  {
    name: "revm",
    paths: ["revm/crates/**/*.rs"],
    extensions: [".rs"],
  },
  {
    name: "execution-specs",
    paths: ["execution-specs/**/*.py"],
    extensions: [".py"],
  },
];

/**
 * Get all Zig code from our EVM implementation
 */
async function getZigCode(): Promise<string> {
  console.log("🔧 Extracting Zig code...");

  try {
    const command = 'find src/evm -name "*.zig"';
    const result = execSync(command, { encoding: "utf8" });
    const zigFiles = result.trim().split("\n").filter(Boolean);

    console.log(`📦 Found ${zigFiles.length} Zig files`);

    const contentParts = [
      "<source>",
      "<description>Our current Zig EVM implementation - the codebase we're working on</description>",
      "",
    ];

    // Read all files concurrently
    const fileReads = zigFiles.map(async (file) => {
      try {
        const content = await readFile(file, "utf8");
        return { file, content };
      } catch (error) {
        console.warn(`⚠️  Could not read ${file}:`, error);
        return null;
      }
    });

    const results = await Promise.all(fileReads);

    for (const result of results) {
      if (result) {
        contentParts.push(
          `## ${result.file}`,
          `\`\`\`zig\n${result.content}\n\`\`\``,
          "",
        );
      }
    }

    contentParts.push("</source>");
    return contentParts.join("\n");
  } catch (error) {
    console.error("❌ Error extracting Zig code:", error);
    return "<source><description>Error extracting Zig code</description></source>";
  }
}

/**
 * Get code from external resource codebase
 */
async function getResourceCode(resource: ResourceConfig): Promise<string> {
  console.log(`🔧 Extracting ${resource.name} code...`);

  try {
    const contentParts = [
      "<source>",
      `<description>${resource.name} reference implementation - external codebase for context</description>`,
      "",
    ];
    let totalFiles = 0;
    const allFileReads: Promise<{
      file: string;
      content: string;
      lang: string;
    } | null>[] = [];

    for (const pattern of resource.paths) {
      const command = `find . -path "./${pattern}" -type f 2>/dev/null`;
      const result = execSync(command, { encoding: "utf8" });
      const files = result.trim().split("\n").filter(Boolean);
      totalFiles += files.length;

      console.log(
        `📦 Found ${files.length} ${resource.name} files for pattern ${pattern}`,
      );

      // Add all file reads to the batch
      for (const file of files) {
        const ext = resource.extensions.find((e) => file.endsWith(e)) || "";
        const lang = ext.replace(".", "") || "text";

        allFileReads.push(
          readFile(file, "utf8")
            .then((content) => ({ file, content, lang }))
            .catch((error) => {
              console.warn(`⚠️  Could not read ${file}:`, error);
              return null;
            }),
        );
      }
    }

    // Read all files concurrently
    const results = await Promise.all(allFileReads);

    // Add successful reads to content
    for (const result of results) {
      if (result) {
        contentParts.push(
          `## ${result.file}`,
          `\`\`\`${result.lang}\n${result.content}\n\`\`\``,
          "",
        );
      }
    }

    contentParts.push("</source>");
    console.log(`📝 Extracted ${totalFiles} files from ${resource.name}`);
    return contentParts.join("\n");
  } catch (error) {
    console.error(`❌ Error extracting ${resource.name} code:`, error);
    return `<source><description>Error extracting ${resource.name} code</description></source>`;
  }
}

/**
 * Create context extraction prompt for a resource
 */
function createContextExtractionPrompt(
  promptContent: string,
  resource: ResourceConfig,
  zigCode: string,
  resourceCode: string,
): string {
  return `You are a prompt engineer helping to add relevant context from ${resource.name} to EVM implementation prompts.

## Your Task
Review the following prompt and the ${resource.name} codebase, then extract ONLY the most relevant code snippets that would help implement the requested feature.

## Original Prompt
${promptContent}

## Our Current Zig EVM Implementation
${zigCode}

## ${resource.name.toUpperCase()} Source Code
${resourceCode}

## Instructions
1. Read the original prompt carefully to understand what feature needs to be implemented
2. Search through the ${resource.name} code to find relevant implementations
3. Extract ONLY the most relevant code snippets that would help with this specific feature
4. Format your response as XML snippets that can be appended to the original prompt
5. Include corrections to the prompt if you notice any inaccuracies
6. Focus on quality over quantity - include only the most helpful context

## Output Format
Provide your response in this exact format:

<${resource.name}>
<file path="github_url_here">
[Include only the most relevant code snippets with proper context]
</file>

<file path="another_github_url_here">
[More relevant code if needed]
</file>
</${resource.name}>

If you found any errors or improvements for the original prompt, include them after the code snippets:

## Prompt Corrections
[Any corrections or improvements to the original prompt]

Remember: Only include code that is directly relevant to implementing the feature described in the prompt. Quality and relevance over quantity.`;
}

/**
 * Process a single prompt file
 */
async function processPromptFile(
  genAI: GoogleGenerativeAI,
  promptFile: string,
  zigCode: string,
  resourceCodes: Map<string, string>,
): Promise<void> {
  const promptName = basename(promptFile, ".md");
  console.log(`\n🔍 Processing prompt: ${promptName}`);

  // Skip WASM-related prompts for now
  if (promptName.toLowerCase().includes("wasm")) {
    console.log(`⏭️  Skipping WASM prompt: ${promptName}`);
    return;
  }

  // Read the original prompt
  const originalPrompt = await readFile(promptFile, "utf8");

  const enhancedParts = [originalPrompt];

  // Process each resource
  for (const resource of RESOURCES) {
    console.log(`  📚 Processing ${resource.name}...`);

    try {
      // Get pre-extracted resource code
      const resourceCode = resourceCodes.get(resource.name) || "";

      // Create context extraction prompt
      const contextPrompt = createContextExtractionPrompt(
        originalPrompt,
        resource,
        zigCode,
        resourceCode,
      );

      console.log(
        `  🤖 Sending to Gemini for ${resource.name} context extraction...`,
      );

      // Send to Gemini
      const model = genAI.getGenerativeModel({ model: MODEL });
      const result = await model.generateContent(contextPrompt);
      const response = result.response;
      const text = response.text();

      console.log(`  ✅ Received response from Gemini for ${resource.name}`);

      // Add to enhanced parts
      enhancedParts.push(`## ${resource.name.toUpperCase()} Context`, text, "");
    } catch (error) {
      console.error(`    ❌ Error processing ${resource.name}:`, error);
    } finally {
      // Add delay to avoid rate limiting (always run whether success or failure)
      await new Promise((resolve) => setTimeout(resolve, 10_000));
    }
  }

  // Write the enhanced prompt
  const enhancedFile = promptFile.replace(".md", "-enhanced.md");
  await writeFile(enhancedFile, enhancedParts.join("\n\n"));

  console.log(`  📝 Enhanced prompt saved to: ${enhancedFile}`);
  console.log(`  ✅ Completed processing: ${promptName}`);
}

/**
 * Main function - processes all prompts
 */
async function main() {
  console.log("🚀 Starting automated prompt enhancement...");

  // Initialize the Gemini API
  const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);

  // Get all prompt files
  const promptsDir = join(process.cwd(), "src/evm/prompts");
  const promptFiles = readdirSync(promptsDir)
    .filter((file) => file.endsWith(".md"))
    .map((file) => join(promptsDir, file));

  console.log(`📋 Found ${promptFiles.length} prompt files to process`);

  // Extract all code once before processing prompts
  console.log("\n📦 Pre-extracting all codebase content...");

  // Extract Zig code once
  const zigCode = await getZigCode();

  // Extract all external resource codes concurrently
  const resourceCodes = new Map<string, string>();
  const resourceExtractions = RESOURCES.map(async (resource) => {
    const code = await getResourceCode(resource);
    resourceCodes.set(resource.name, code);
  });

  await Promise.all(resourceExtractions);

  console.log("✅ All codebase content extracted and cached");

  // Process only the first 1 file for testing
  const filesToProcess = promptFiles.slice(0, 1);
  console.log(
    `\n🔍 Processing first ${filesToProcess.length} file for testing:`,
  );
  filesToProcess.forEach((file, index) => {
    console.log(`  ${index + 1}. ${basename(file)}`);
  });

  // Process each prompt file
  for (const promptFile of filesToProcess) {
    await processPromptFile(genAI, promptFile, zigCode, resourceCodes);
  }

  console.log("\n🎉 All prompts processed successfully!");
}

main().catch(console.error);