// @category _Reman3
// @menupath Reman3.Rebuild Function Database

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import re3lib.FunctionDatabase;
import re3lib.RecompileConfig;

import java.io.File;
import java.io.BufferedReader;
import java.io.FileReader;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class RebuildFunctionDatabase extends GhidraScript {
  FunctionDatabase functionDB;

  @Override
  public void run() throws Exception {
    RecompileConfig.INSTANCE = new RecompileConfig(this);

    functionDB = new FunctionDatabase(this);

    scanDirectory(RecompileConfig.INSTANCE.dirDecompAuto);
    scanDirectory(RecompileConfig.INSTANCE.dirDecompFix);
    scanDirectory(RecompileConfig.INSTANCE.dirDecompStub);

    println("Applying default filters...");
    functionDB.applyDefaultFilters();

    println("Saving function database...");
    functionDB.save();

    println("Function database rebuilt successfully.");

    // for (FunctionDatabase.Entry entry : functionDB.entries) {
    //   println(entry.address + " " + entry.name + " " + entry.file.getName());
    //   for (FunctionDatabase.Dependency dependency : entry.dependencies) {
    //     println("  " + dependency.address + " " + dependency.name);
    //   }
    // }
  }

  private void scanDirectory(File directory) throws Exception {
    File[] files = directory.listFiles((dir, name) -> name.endsWith(".cxx"));
    if (files == null)
      return;

    for (File file : files) {
      scanFile(file);
    }
  }

  private void scanFile(File file) throws Exception {
    println("Scanning " + file);
    try (BufferedReader reader = new BufferedReader(new FileReader(file))) {
      String line;
      Pattern dependencyPattern = Pattern.compile("(\\w+)\\s+(\\w+)\\(.*\\);\\s*//\\s*([0-9A-Fa-f]{8})\\s*//\\s*(.*)");
      Pattern addressPattern = Pattern.compile("//\\s*([0-9A-Fa-f]{8})");
      Pattern functionNamePattern = Pattern.compile("(\\S+)\\s+(\\S+)\\s*\\(");

      List<FunctionDatabase.Dependency> dependencies = new ArrayList<>();
      String address = null;
      String functionName = null;

      while ((line = reader.readLine()) != null) {
        Matcher dependencyMatcher = dependencyPattern.matcher(line);
        if (dependencyMatcher.find()) {
          // println("Found dependency: " + dependencyMatcher.group(3));
          Address depAddress = currentProgram.getAddressFactory().getAddress(dependencyMatcher.group(3));
          String name = dependencyMatcher.group(2);
          FunctionDatabase.Dependency dependency = functionDB.new Dependency(depAddress, name);
          dependencies.add(dependency);
          continue;
        }

        Matcher addressMatcher = addressPattern.matcher(line);
        if (addressMatcher.find()) {
          // println("Found address: " + addressMatcher.group(1));
          address = addressMatcher.group(1);
          // Skip any comments or newlines between address and function definition
          while ((line = reader.readLine()) != null) {
            line = line.trim();
            // println("Line: " + line);
            if (!line.isEmpty()) {
              Matcher functionNameMatcher = functionNamePattern.matcher(line);
              if (functionNameMatcher.find()) {
                functionName = functionNameMatcher.group(2).trim();
                break;
              }
            }
          }
          if (functionName != null) {
            break;
          }
        }
      }

      if (address != null && functionName != null) {
        Address functionAddress = currentProgram.getAddressFactory().getAddress(address);
        FunctionDatabase.Entry entry = functionDB.new Entry();
        entry.address = functionAddress;
        entry.name = functionName;
        entry.file = file;
        entry.dependencies = dependencies;
        functionDB.entries.add(entry);
      } else {
        throw new Exception("Failed to parse function at " + file.getName());
      }
    }
  }
}