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

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.StandAloneDataTypeManager;
import re3lib.FunctionDatabase;
import re3lib.RecompileConfig;
import re3lib.CParser;
import re3lib.CTokenizer;

import java.io.File;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class RebuildFunctionDatabase extends GhidraScript {
  // Will rebuild all functions
  public boolean rebuildAllGlobals = true;
  FunctionDatabase functionDB;

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

    functionDB = new FunctionDatabase(this);

    scanFile(new File(RecompileConfig.INSTANCE.outputDir, "gh_auto/r3_engineLoop.cxx"), FunctionDatabase.Type.Auto);
    // scanDirectory(RecompileConfig.INSTANCE.dirDecompAuto, FunctionDatabase.Type.Auto);
    // scanDirectory(RecompileConfig.INSTANCE.dirDecompFix, FunctionDatabase.Type.Fix);
    // scanDirectory(RecompileConfig.INSTANCE.dirDecompStub, FunctionDatabase.Type.Stub);

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

    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, FunctionDatabase.Type type) throws Exception {
    File[] files = directory.listFiles((dir, name) -> name.endsWith(".cxx"));
    if (files == null)
      return;

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

  private void parseOld(BufferedReader reader, File file, FunctionDatabase.Type type) throws Exception {
    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.type = type;
      entry.dependencies = dependencies;
      functionDB.entries.add(entry);
    } else {
      // throw new Exception("Failed to parse function at " + file.getName());
      println("Failed to parse function at " + file.getName());
    }
  }

  private void scanFile(File file, FunctionDatabase.Type type) throws Exception {
    println("Scanning " + file);

    String text = new String(Files.readAllBytes(file.toPath()));
    CTokenizer.TokenSet tokens = new CTokenizer(text).parse();
    CParser parser = new CParser(tokens);
    parser.parse();

    // for (CTokenizer.Token token : tokens.getTokens()) {
    //   int line = tokens.getLine(token.ofs);
    //   println("Line " + line + ": " + token.ofs + " " + token.len + " " + token.type + " - "
    //       + tokens.getTextNoNewlines(token));
    // }
    for (CParser.Function function : parser.getFunctions()) {
      println("Function: " + function.name + " " + function.startOffset + " " + function.endOffset);
    }
    for (CParser.FunctionCall functionCall : parser.getFunctionCalls()) {
      println("FunctionCall: " + functionCall.name + " " + functionCall.startOffset + " " + functionCall.endOffset);
    }
    for (CParser.Variable variable : parser.getVariables()) {
      println("Variable: " + variable.name + " " + variable.startOffset + " " + variable.endOffset);
    }
  }
}