package re3lib;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Dictionary;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;

import javax.xml.datatype.DatatypeFactory;

import ghidra.app.cmd.label.AddLabelCmd;
import ghidra.app.script.GhidraScript;
import ghidra.app.services.DataTypeManagerService;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.AbstractStringDataType;
import ghidra.program.model.data.Array;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.DataTypeManager;
import ghidra.program.model.data.DataTypePath;
import ghidra.program.model.data.PointerDataType;
import ghidra.program.model.data.Undefined;
import ghidra.program.model.listing.Data;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.Symbol;
import ghidra.util.data.DataTypeParser;
import ghidra.util.data.DataTypeParser.AllowedDataTypes;

public class GlobalDumper {

  public class GlobalRec {
    public Address address;
    public String name;
    public DataType type;

    public GlobalRec(Address address, String name, DataType type) {
      this.address = address;
      this.name = name;
      this.type = type;
    }
  };

  GhidraScript script;
  File manifestFile;
  HashMap<Address, GlobalRec> globalAddrs = new HashMap<>();

  public GlobalDumper(GhidraScript script) {
    this.script = script;
    manifestFile = new File(RecompileConfig.INSTANCE.outputDir, "globals.txt");
  }

  public boolean loadGlobalManifest() throws Exception {
    // Globals are stored in the format of
    // <address> || <name> || <type>

    if (!manifestFile.exists()) {
      script.println("Global manifest file not found: " + manifestFile);
      return false;
    }

    // Get the dataTypeManagerService
    DataTypeManagerService dataTypeManagerService = (DataTypeManagerService) script.getState().getTool()
        .getService(DataTypeManagerService.class);
    DataTypeManager dtm = script.getCurrentProgram().getDataTypeManager();
    DataTypeParser dtp = new DataTypeParser(dataTypeManagerService, AllowedDataTypes.ALL);
    try (BufferedReader reader = new BufferedReader(new FileReader(manifestFile))) {
      String line;
      while ((line = reader.readLine()) != null) {
        String[] parts = line.split("\\|\\|");
        if (parts.length == 4) {
          Address address = script.parseAddress(parts[0].trim());
          String name = parts[1].trim();
          String categoryPath = parts[2].trim();
          String dataTypePath = parts[3].trim();
          DataTypePath typePath = new DataTypePath(categoryPath, dataTypePath);
          DataType type = null;
          type = dtm.getDataType(typePath);
          if (type == null) {
            // script.println("Parsing type: " + dataTypePath);
            type = dtp.parse(dataTypePath);
          }
          if (type == null) {
            script.println("WARNING: Failed to find type: " + dataTypePath + " for global: " + name + " at " + address);
            continue;
          }
          globalAddrs.put(address, new GlobalRec(address, name, type));
        } else {
          script.println("Invalid global manifest line: " + line);
        }
      }
    }

    script.println("Loaded " + globalAddrs.size() + " globals from " + manifestFile);
    return true;
  }

  public void addGlobal(Address addr, HighSymbol sym) {
    DataType dt = sym.getDataType();
    if (dt == null) {
      script.println("WARNING: Missing type for global: " + sym.getName() + " at " + addr);
      return;
    }
    globalAddrs.put(addr, new GlobalRec(addr, sym.getName(), dt));
  }

  String escapeCString(String str) {
    str = str.replace("\\", "\\\\");
    str = str.replace("\"", "\\\"");
    return str;
  }

  String readCString(Address addr, int maxLen) throws Exception {
    StringBuilder sb = new StringBuilder();
    int ofs = 0;
    while (true) {
      Address read = addr.add(ofs++);
      // println("Reading: " + read);
      byte b = script.getCurrentProgram().getMemory().getByte(read);
      // println("Read: " + b);
      if (b == 0 || ofs >= maxLen) {
        break;
      }
      sb.append((char) b);
    }
    if (sb.length() > 0) {
      // println("STR \"" + sb.toString() + "\"");
    }
    return sb.toString();
  }

  public void dumpGlobals() throws Exception {
    File globalSymbolsListH = new File(RecompileConfig.INSTANCE.outputDir, "gh_global.h");
    PrintWriter hwriter = new PrintWriter(globalSymbolsListH, "UTF-8");
    hwriter.println("// AUTO-GENERATED FILE ");
    Utils.headerGuardPre(hwriter, "GLOBALS");
    hwriter.println("#include <gh_global_binder.h>");
    hwriter.println();

    File globalSymbolsListC = new File(RecompileConfig.INSTANCE.outputDir, "gh_global.cxx");
    PrintWriter cwriter = new PrintWriter(globalSymbolsListC, "UTF-8");
    cwriter.println("// AUTO-GENERATED FILE ");
    cwriter.println("#include <gh_global_binder.h>");
    hwriter.println();

    List<GlobalRec> globals = new ArrayList<>(globalAddrs.values());
    globals.sort((o1, o2) -> o1.address.compareTo(o2.address));

    for (GlobalRec global : globals) {
      DataType dt = global.type;
      if (dt == null) {
        script.println("WARNING: Missing type for global: " + global.name + " at " + global.address);
        continue;
      }
      String dataType = dt.getDisplayName();
      String sanitizedName = Utils.sanitizeIdentifier(global.name);
      String name = global.name;
      Address addr = global.address;
      if (!sanitizedName.equals(name)) {
        script.println("WARNING: Invalid global symbol name: " + name);
        name = sanitizedName;
      }
      // println("Symbol: " + symbol + " Addr: " + addr + " Size:" + symSize + " " +
      // storage.getSerializationString());
      try {
        String initBlk = " = ";
        boolean fullyDefinedType = false;
        if (dt instanceof AbstractStringDataType) {
          AbstractStringDataType sdt = (AbstractStringDataType) dt;
          dataType = "const char*";
          // String type
          initBlk += "\"" + escapeCString(readCString(addr, 2048)) + "\"";
          fullyDefinedType = true;
        } else if (dt instanceof PointerDataType) {
          PointerDataType pdt = (PointerDataType) dt;
          DataType baseType = pdt.getDataType();
          dataType = baseType.getDisplayName() + "*";
          initBlk += "(" + dataType + ")&GH_MEM(0x" + addr + ")";
          fullyDefinedType = true;
        }
        String linkagePrefix = "extern ";
        if (fullyDefinedType) {
          hwriter.println(linkagePrefix + dataType + " " + name + "; // " + addr);
          cwriter.println(dataType + " " + name + initBlk + "; // " + addr);
        } else {
          if (dt instanceof Array) {
            // println("Array: " + dt.getDisplayName() + " - " + addr + " - " +
            // dt.getClass().getSimpleName());
            Array adt = (Array) dt;
            DataType baseType = adt.getDataType();
            hwriter.println(
              linkagePrefix + baseType.getDisplayName() + "(&" + name + ")[" + adt.getNumElements() + "]; // " + addr);
            cwriter.println(
                baseType.getDisplayName() + "(&" + name + ")[" + adt.getNumElements() + "] = *reinterpret_cast<"
                    + baseType.getDisplayName() + "(*)[" + adt.getNumElements() + "]>(GH_MEM(0x" + addr + "));");
          } else {
            String refTypeStr = dt.getDisplayName() + "&";
            hwriter.println(linkagePrefix + refTypeStr + " " + name + "; // " + addr);
            cwriter.println(refTypeStr + " " + name + "= (" + refTypeStr + ") GH_MEM(0x" + addr + ");");
          }
        }
      } catch (Exception e) {
        script.println("Error processing global symbol: " + e);
        script.println("Symbol: " + name + " - " + addr);
      }
    }

    Utils.headerGuardPost(hwriter, "GLOBALS");
    hwriter.close();
    cwriter.close();
  }

  public void saveGlobalManifest() throws Exception {
    File backupFile = new File(manifestFile.getParentFile(), manifestFile.getName() + ".bak");
    if (backupFile.exists()) {
      if (!backupFile.delete()) {
        throw new Exception("Failed to delete backup file: " + backupFile + ", globals will not be saved!");
      }
    }
    if (!manifestFile.renameTo(backupFile))
      throw new Exception("Failed to rename manifest file: " + manifestFile + ", globals will not be saved!");

    try (PrintWriter writer = new PrintWriter(manifestFile)) {
      script.println("Saving global manifest to " + manifestFile);
      GlobalRec[] globals = globalAddrs.values().toArray(new GlobalRec[0]);
      Arrays.sort(globals, (a, b) -> a.address.compareTo(b.address));
      for (GlobalRec global : globals) {
        DataTypePath path = global.type.getDataTypePath();
        writer.println(global.address.toString() + " || " + global.name + " || " + path.getCategoryPath() + " || "
            + path.getDataTypeName());
      }
    }
  }

  public void sanitizeGlobalSymbols() {
    for (GlobalRec global : globalAddrs.values()) {
      String sanitizedName = Utils.sanitizeIdentifier(global.name);
      if (!sanitizedName.equals(global.name)) {
        Symbol symbol = script.getSymbolAt(global.address);
        if (symbol != null) {
          script.println("Renaming global symbol: " + global.name + " -> " + sanitizedName);
          AddLabelCmd cmd = new AddLabelCmd(global.address, sanitizedName,
              symbol.getParentNamespace(),
              SourceType.USER_DEFINED);
          if (!cmd.applyTo(script.getCurrentProgram())) {
            script.println("Error renaming symbol: " + cmd.getStatusMsg());
          }
          global.name = sanitizedName;
        }
      }
    }
  }
}