// Cleanup the database of all missing files, and report duplicates
// @category _Reman3
// @menupath Reman3.Cleanup Database  

import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.Function;
import re3lib.FunctionDumper;
import re3lib.GlobalDumper;
import re3lib.RemanConfig;
import re3lib.FunctionDatabase;
import java.util.List;
import java.util.Map;
import java.util.HashMap;
import java.util.ArrayList;
import ghidra.program.model.address.Address;

public class CleanupDatabase extends GhidraScript {
  @Override
  public void run() throws Exception {
    RemanConfig.INSTANCE = new RemanConfig(this);
    RemanConfig.INSTANCE.createDirectories();
    boolean anyChanges = false;

    try (FunctionDatabase functionDatabase = new FunctionDatabase(this)) {
      functionDatabase.connect();

      // Step 1: Check for missing files and remove entries
      anyChanges |= cleanupMissingFiles(functionDatabase);

      // Step 2: Check for duplicate addresses
      anyChanges |= reportDuplicateAddresses(functionDatabase);

      // Step 3: Check if database names match function names
      anyChanges |= validateFunctionNames(functionDatabase);

      if (anyChanges) {
        println("Database cleanup completed with changes - touching CMake timestamp");
        RemanConfig.INSTANCE.touchCMakeTimestamp();
      } else {
        println("Database cleanup completed - no changes needed");
      }
    }
  }

  private boolean cleanupMissingFiles(FunctionDatabase functionDatabase) throws Exception {
    println("=== Step 1: Checking for missing files ===");
    boolean madeChanges = false;

    List<FunctionDatabase.FunctionEntry> entries = functionDatabase.loadAllEntries();
    List<FunctionDatabase.FunctionEntry> entriesToRemove = new ArrayList<>();

    for (FunctionDatabase.FunctionEntry entry : entries) {
      if (!entry.file.exists()) {
        println("Missing file: " + entry.file.getAbsolutePath() + " for function " + entry.name);
        entriesToRemove.add(entry);
      }
    }

    for (FunctionDatabase.FunctionEntry entry : entriesToRemove) {
      println("Removing database entry for missing file: " + entry.file.getAbsolutePath());
      functionDatabase.removeEntryAt(entry.file.getAbsolutePath());
      madeChanges = true;
    }

    println("Removed " + entriesToRemove.size() + " entries for missing files");
    return madeChanges;
  }

  private boolean reportDuplicateAddresses(FunctionDatabase functionDatabase) throws Exception {
    println("=== Step 2: Checking for duplicate addresses ===");
    boolean madeChanges = false;

    List<FunctionDatabase.FunctionEntry> entries = functionDatabase.loadAllEntries();
    Map<Address, List<FunctionDatabase.FunctionEntry>> addressMap = new HashMap<>();

    // Group entries by address
    for (FunctionDatabase.FunctionEntry entry : entries) {
      if (entry.type == FunctionDatabase.Type.Ref)
        continue;
      if (!addressMap.containsKey(entry.address)) {
        addressMap.put(entry.address, new ArrayList<>());
      }
      addressMap.get(entry.address).add(entry);
    }

    // Check for duplicates
    for (Map.Entry<Address, List<FunctionDatabase.FunctionEntry>> mapEntry : addressMap.entrySet()) {
      List<FunctionDatabase.FunctionEntry> duplicates = mapEntry.getValue();
      if (duplicates.size() > 1) {
        println("Found " + duplicates.size() + " entries for address " + mapEntry.getKey() + ":");

        // Report all entries
        for (FunctionDatabase.FunctionEntry entry : duplicates) {
          println("  - " + entry.type + ": " + entry.file.getAbsolutePath() + " (name: " + entry.name + ")");
        }
      }
    }

    return madeChanges;
  }

  private boolean validateFunctionNames(FunctionDatabase functionDatabase) throws Exception {
    println("=== Step 3: Validating function names ===");
    boolean madeChanges = false;

    List<FunctionDatabase.FunctionEntry> entries = functionDatabase.loadAllEntries();
    List<FunctionDatabase.FunctionEntry> entriesToUpdate = new ArrayList<>();
    List<FunctionDatabase.FunctionEntry> entriesToRemove = new ArrayList<>();

    for (FunctionDatabase.FunctionEntry entry : entries) {
      Function function = getFunctionAt(entry.address);

      if (function == null) {
        println("No function found at address " + entry.address + " for database entry: " + entry.name);
        println("  Removing orphaned database entry");
        entriesToRemove.add(entry);
      } else {
        String actualName = function.getName();
        if (!actualName.equals(entry.name)) {
          println("Name mismatch at " + entry.address + ":");
          println("  Database name: " + entry.name);
          println("  Actual name: " + actualName);

          // Update the entry with the correct name
          entry.name = actualName;
          entriesToUpdate.add(entry);
        }
      }
    }

    // Report orphaned entries
    for (FunctionDatabase.FunctionEntry entry : entriesToRemove) {
      println("Found orphaned entry: " + entry.file.getAbsolutePath() + " for function " + entry.name + " at address "
          + entry.address);
    }

    // Update entries with corrected names
    for (FunctionDatabase.FunctionEntry entry : entriesToUpdate) {
      println("Updating database entry for " + entry.address + " with correct name: " + entry.name);
      functionDatabase.addEntryAt(entry);
      madeChanges = true;
    }

    println("Removed " + entriesToRemove.size() + " orphaned entries");
    println("Updated " + entriesToUpdate.size() + " entries with corrected names");

    return madeChanges;
  }
}