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.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;
  FunctionDatabase functionDatabase;
  HashMap<Address, GlobalRec> globalAddrs = new HashMap<>();

  public GlobalDumper(GhidraScript script, FunctionDatabase functionDatabase) {
    this.script = script;
    this.functionDatabase = functionDatabase;
  }

  public void removeGlobalManifest() {
    // Remove globals from database instead of file
    try {
      functionDatabase.removeGlobalsByFilepath(RemanConfig.INSTANCE.GLOBAL_H_FILE);
      globalAddrs.clear();
    } catch (Exception e) {
      script.println("Error removing global manifest: " + e.getMessage());
    }
  }

  public boolean loadGlobalManifest() throws Exception {
    globalAddrs.clear();
    
    // Load globals from database
    List<FunctionDatabase.GlobalEntry> dbGlobals = functionDatabase.loadAllGlobals();
    
    // Get the dataTypeManagerService for parsing types
    DataTypeManagerService dataTypeManagerService = (DataTypeManagerService) script.getState().getTool()
        .getService(DataTypeManagerService.class);
    DataTypeManager dtm = script.getCurrentProgram().getDataTypeManager();
    DataTypeParser dtp = new DataTypeParser(dataTypeManagerService, AllowedDataTypes.ALL);
    
    for (FunctionDatabase.GlobalEntry entry : dbGlobals) {
      // Note: The database stores type as string, need to reconstruct DataType
      // For now, we'll parse it back from the type string stored in database
      // This is a limitation of moving from the manifest format
      DataType type = null;
      
      // Try to get from existing data at address
      Data data = script.getDataAt(entry.address);
      if (data != null) {
        type = data.getDataType();
      }
      
      if (type == null) {
        script.println("WARNING: Could not reconstruct type for global: " + entry.name + " at " + entry.address);
        type = dtp.parse("undefined4");
      }

      // Retrieve the name again
      Symbol sym = script.getSymbolAt(entry.address);
      if (sym != null) {
        entry.name = sym.getName();
      }
      
      globalAddrs.put(entry.address, new GlobalRec(entry.address, entry.name, type));
    }

    script.println("Loaded " + globalAddrs.size() + " globals from database");
    return !globalAddrs.isEmpty();
  }

  String escapeCString(String str) {
    str = str.replace("\\", "\\\\");
    str = str.replace("\"", "\\\"");
    str = str.replace("\n", "\\n");
    str = str.replace("\r", "\\r");
    str = str.replace("\t", "\\t");
    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(RemanConfig.INSTANCE.outputDir, RemanConfig.GLOBAL_H_FILE);
    PrintWriter hwriter = new PrintWriter(globalSymbolsListH, "UTF-8");
    hwriter.println("// AUTO-GENERATED FILE ");
    Utils.headerGuardPre(hwriter, "GLOBALS");
    hwriter.println("#include <r3/binders/global.h>");
    hwriter.println();

    File globalSymbolsListC = new File(RemanConfig.INSTANCE.outputDir, RemanConfig.GLOBAL_CXX_FILE);
    PrintWriter cwriter = new PrintWriter(globalSymbolsListC, "UTF-8");
    cwriter.println("// AUTO-GENERATED FILE ");
    cwriter.println("#include <r3/binders/global.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 {
    // Save globals to database instead of file
    script.println("Saving globals to database");
    
    // Clear existing globals for the default filepath
    functionDatabase.removeGlobalsByFilepath(RemanConfig.GLOBAL_H_FILE);
    
    // Add all current globals to database
    for (GlobalRec global : globalAddrs.values()) {
      String dataTypeName = global.type.getDisplayName();
      functionDatabase.addGlobal(global.address, global.name, dataTypeName);
    }
  }

  public void addGlobal(Address addr, HighSymbol sym) throws Exception {
    if (sym.getName().equals("ExceptionList")) {
      return;
    }

    if (addr.compareTo(RemanConfig.INSTANCE.staticMemoryBlockStart) < 0
        || addr.compareTo(RemanConfig.INSTANCE.staticMemoryBlockEnd) > 0) {
      throw new Exception("Global address out of range: " + addr);
    }

    DataType dt = sym.getDataType();
    // if(symb.get
    if (sym.getDataType().getName() == "undefined") {
      // script.println("UNDEFINED: " + addr + " - " + dt.getDisplayName() + " - " +
      // dt.getClass().getName());
      Data data = script.getDataAt(addr);
      if (data != null) {
        dt = data.getDataType();
        // script.println("DATA: " + addr + " - " + dt.getDisplayName());
      }
    }
    if (dt == null) {
      script.println("WARNING: Missing type for global: " + sym.getName() + " at " + addr);
      return;
    }
    // script.println("Global: " + addr + " - " + sym.getName() + " - " +
    // dt.getDisplayName());
    globalAddrs.put(addr, new GlobalRec(addr, sym.getName(), dt));
  }

  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;
        }
      }
    }
  }
}