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 ghidra.app.script.GhidraScript;
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.PointerDataType;
import ghidra.program.model.pcode.HighSymbol;

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 void loadGlobalManifest() throws Exception {
    // Globals are stored in the format of
    // <address> || <name> || <type>
    try (BufferedReader reader = new BufferedReader(new FileReader(manifestFile))) {
      String line;
      while ((line = reader.readLine()) != null) {
        String[] parts = line.split(" || ");
        if (parts.length == 3) {
          String address = parts[0];
          String name = parts[1];
          String typeName = parts[2];
          DataType type = script.getCurrentProgram().getDataTypeManager().getDataType(typeName);
          globalAddrs.put(script.parseAddress(address), new GlobalRec(script.parseAddress(address), name, type));
        }
      }
    }
  }

  public void addGlobal(Address addr, HighSymbol sym) {
    globalAddrs.put(addr, new GlobalRec(addr, sym.getName(), sym.getDataType()));
  }

  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;
        }
        if (fullyDefinedType) {
          hwriter.println("extern " + 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(
                "extern " + 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("extern " + refTypeStr + " " + name + "; // " + addr);
            cwriter.println(dataType + " " + 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 {
    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) {
        writer.println(global.address.toString() + " || " + global.name + " || " + global.type.getDisplayName());
      }
    }
  }
}