package re3lib;

import java.io.File;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import ghidra.app.decompiler.ClangLine;
import ghidra.app.decompiler.ClangSyntaxToken;
import ghidra.app.decompiler.ClangToken;
import ghidra.app.decompiler.DecompileResults;
import ghidra.app.decompiler.PrettyPrinter;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.pcode.HighFunction;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.Varnode;

public class FunctionDumper {
  GhidraScript script;
  GlobalDumper globalDumper;
  FunctionDatabase functionDatabase;

  public HashSet<Address> functionAddrBlackList = new HashSet<>();

  public boolean createdFile = false;
  // Collects functions called by the current function
  public HashSet<Function> functionReferences = new HashSet<>();

  static final Pattern fieldAccessRegex = Pattern.compile("^_([0-9]+)_([0-9]+)_$");

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

  boolean isValidFunction(Function function) {
    if (functionAddrBlackList.contains(function.getEntryPoint()))
      return false;
    if (function.getComment() != null) {
      if (function.getComment().startsWith("/i"))
        return false;
      if (function.getComment().startsWith("library function"))
        return false;
    }
    if (function.getName().startsWith("crt_"))
      return false;
    return true;
  }

  void initFunctionBlacklist() {
    functionAddrBlackList = Utils.loadFunctionBlacklist(RemanConfig.INSTANCE.functionBlacklistPath);

    // Build blacklist if not loaded
    if (functionAddrBlackList == null) {
      boolean modified = false;
      Iterator<Function> functionsIt = script.getCurrentProgram().getFunctionManager().getFunctions(true).iterator();
      while (functionsIt.hasNext()) {
        Function function = functionsIt.next();
        if (functionAddrBlackList.contains(function.getEntryPoint())) {
          continue;
        }

        String comment = function.getComment();
        boolean isIgnoredFunction = false;
        if (comment != null && comment.contains("Library Function")) {
          script.println("Adding library function " + function.getName() + " to blacklist");
          script.println("ac:" + functionAddrBlackList.size() + " jj:"
              + functionAddrBlackList.contains(function.getEntryPoint()) + " "
              + function.getEntryPoint());
          isIgnoredFunction = true;
        }

        if (function.getName().startsWith("crt_")) {
          script.println("Adding crt function " + function.getName() + " to blacklist");
          isIgnoredFunction = true;
        }

        if (isIgnoredFunction) {
          // Decompile and trace
          PCallTracer tracer = new PCallTracer();
          tracer.setBlacklist(functionAddrBlackList);
          tracer.traceCalls(function);
          for (Function f : tracer.out) {
            script.println(" Adding " + f.getName() + " to blacklist");
            functionAddrBlackList.add(f.getEntryPoint());
            modified = true;
          }
        }
      }

      if (modified) {
        Utils.saveFunctionBlacklist(functionAddrBlackList, RemanConfig.INSTANCE.functionBlacklistPath);
      }
    }
  }

  public static boolean isDumpedFix(Function function) {
    String sanitizedFunctionName = Utils.sanitizeIdentifier(function.getName());
    String fileName = sanitizedFunctionName + ".cxx";
    File f0 = new File(RemanConfig.INSTANCE.dirDecompFix, fileName);
    return f0.exists();
  }

  public static boolean isDumpedAuto(Function function) {
    String sanitizedFunctionName = Utils.sanitizeIdentifier(function.getName());
    String fileName = sanitizedFunctionName + ".cxx";
    File f0 = new File(RemanConfig.INSTANCE.dirDecompAuto, fileName);
    return f0.exists();
  }

  public void dump(Function function)
      throws Exception {
    String sanitizedFunctionName = Utils.sanitizeIdentifier(function.getName());
    String fileName = sanitizedFunctionName + ".cxx";

    Address entrypoint = function.getEntryPoint();
    List<FunctionDatabase.FunctionEntry> entries = functionDatabase.findEntriesByAddress(entrypoint);
    FunctionDatabase.Type targetType = FunctionDatabase.Type.Auto;
    for (FunctionDatabase.FunctionEntry entry : entries) {
      script.println("Found existing decompiled entry at " + entry.file + " - " + entry.name);
      if (targetType != FunctionDatabase.Type.Ref) {
        if (entry.type == FunctionDatabase.Type.Fix) {
          targetType = FunctionDatabase.Type.Ref;
        }
      }
      if (entry.type == FunctionDatabase.Type.Stub) {
        // Remove the stub file, since we now use the decompiled file
        File stubFile = entry.file;
        if (stubFile.exists()) {
          script.println("Removing function stub " + stubFile);
          stubFile.delete();
          createdFile = true;
          functionDatabase.removeEntryAt(entry.file.toString());
        }
      }
    }

    File targetFilename = null;
    if (targetType == FunctionDatabase.Type.Ref) {
      targetFilename = new File(RemanConfig.INSTANCE.dirDecompRef, fileName);
    } else {
      targetFilename = new File(RemanConfig.INSTANCE.dirDecompAuto, fileName);
    }
    if (targetFilename.exists()) {
      targetFilename.delete();
      script.println("Overwriting existing file " + targetFilename);
    } else {
      createdFile = true;
    }

    File f0 = targetFilename;
    script.println("Processing " + function.getName() + " => " + f0.toString());

    // Update database
    FunctionDatabase.FunctionEntry newEntry = new FunctionDatabase.FunctionEntry(entrypoint, function.getName(), f0,
        targetType);
    functionDatabase.addEntryAt(newEntry);

    List<Function> externalFunctionCalls = new ArrayList<>();

    DecompileResults decompRes = RemanConfig.INSTANCE.decompCache.getOrInsert(function);
    try (PrintWriter writer2 = new PrintWriter(f0, "UTF-8")) {
      writer2.println("// AUTO-GENERATED FILE, MOVE TO 'gh_fix' FOLDER PREVENT OVERWRITING!!!!! ");
      writer2.println();
      writer2.println("#include <r3/binders/auto.h>");
      writer2.println("#include <gh_global.h>");
      writer2.println();
      writer2.println("extern \"C\" {");

      HighFunction highFunction = decompRes.getHighFunction();

      HashSet<String> headers = new HashSet<>();
      StringWriter codeWriter = new StringWriter();

      PrettyPrinter pp = new PrettyPrinter(decompRes.getFunction(), decompRes.getCCodeMarkup(), null);
      Iterator<ClangLine> lines = pp.getLines().iterator();
      while (lines.hasNext()) {
        ClangLine line = lines.next();
        for (int i = 0; i < line.getIndent(); i++) {
          codeWriter.write(' ');
        }

        List<ClangToken> tokens = new ArrayList<>();
        // Parse preliminary line tokens
        for (int i = 0; i < line.getNumTokens(); i++) {
          ClangToken token = line.getToken(i);
          if (token.getText().equals("__cdecl") || token.getText().equals("__thiscall")
              || token.getText().equals("__stdcall")) {
            // Remove function declaration
            continue;
          }
          if (!token.getText().isEmpty())
            tokens.add(token);
        }

        // Preprocess tokens
        boolean prevDot = false;
        for (int t = 0; t < tokens.size(); t++) {
          ClangToken token = tokens.get(t);

          boolean thisDot = false;
          // script.println("Token: " + token.toString());
          if (token.toString().equals(".")) {
            // println("Found dot: " + token.toString() + " - " + token.getClass());
            thisDot = true;
          }

          if (prevDot) {
            // println("Possible field access: " + token.getText());
            if (token instanceof ClangSyntaxToken) {
              // Parse _4_4_ sub-access using regex
              String text = token.getText();
              Matcher matcher = fieldAccessRegex.matcher(text);
              if (matcher.matches()) {
                int offset = Integer.parseInt(matcher.group(1));
                int size = Integer.parseInt(matcher.group(2));
                // println("MATCHED: " + token.getText() + " - " + token.getSyntaxType() + " - "
                // + token.getVarnode() + " - "
                // + token.getPcodeOp());

                // Replace tokens with + Field<offset, size>
                ClangToken replacement = new ClangToken(token.Parent(), " + Field<" + offset + ", " + size + ">()");
                tokens.remove(t);
                tokens.remove(t - 1);
                tokens.add(t - 1, replacement);
                t--;
              }
            }
          }

          // Extract memory references
          HighSymbol gsym = token.getHighSymbol(highFunction);
          if (gsym != null) {
            var symStorage = gsym.getStorage();
            var sym = gsym.getSymbol();

            Address address;
            if (symStorage.isUnassignedStorage()) {
              address = sym.getAddress();
            } else {
              address = gsym.getStorage().getMinAddress();
            }

            // Check if it's a function pointer, otherwise add to globals
            if (address.isMemoryAddress()) {
              // script.println("Address: " + address + " - " + sym.getName());
              Function maybeFunction = script.getFunctionAt(address);
              if (maybeFunction != null) {
                externalFunctionCalls.add(maybeFunction);
              } else {
                script.println("Add globals " + address + " - " + gsym.getName() + " - " + token.getText());
                globalDumper.addGlobal(address, gsym);
              }
            }
          }

          // Extract external function calls
          PcodeOp op = token.getPcodeOp();
          if (op != null && op.getOpcode() == PcodeOp.CALL) {
            // println("PcodeOp: " + op.toString() + " - " + op.getInput(0).toString());
            Varnode target = op.getInput(0);
            if (target.isAddress()) {
              Address callAddr = target.getAddress();
              Function calledFunction = script.getFunctionAt(callAddr);
              if (calledFunction != null) {
                if (isValidFunction(calledFunction)) {
                  externalFunctionCalls.add(calledFunction);
                }
              }
            }
          }
          prevDot = thisDot;
        }

        // Print tokens
        for (int t = 0; t < tokens.size(); t++) {
          ClangToken token = tokens.get(t);
          codeWriter.write(token.toString());
        }
        codeWriter.write('\n');
      }

      for (Function externalFunction : externalFunctionCalls) {
        String proto = externalFunction.getSignature().getPrototypeString(false);
        headers.add("" + proto
            + "; // " + externalFunction.getEntryPoint() + " // "
            + externalFunction.getName());
      }

      for (String header : headers) {
        writer2.println(header);
      }
      writer2.println();
      writer2.print("// " + function.getEntryPoint());
      writer2.print(codeWriter.toString());
      writer2.println("}");
      writer2.println();
    }

    // Possibly generate stubs for external functions
    for (Function externalFunction : externalFunctionCalls) {
      String sanitizedExtFunctionName = Utils.sanitizeIdentifier(externalFunction.getName());

      List<FunctionDatabase.FunctionEntry> entries1 = functionDatabase
          .findEntriesByAddress(externalFunction.getEntryPoint());
      boolean needStub = true;
      for (FunctionDatabase.FunctionEntry entry : entries1) {
        if (entry.type == FunctionDatabase.Type.Auto || entry.type == FunctionDatabase.Type.Fix) {
          needStub = false;
          break;
        }
      }
      if (!needStub)
        continue;

      fileName = sanitizedExtFunctionName + ".cxx";
      File f4 = new File(RemanConfig.INSTANCE.dirDecompStub, fileName);
      script.println("Generating function stub for " + externalFunction.getName() + " => " + f4.toString());

      try (PrintWriter writer2 = new PrintWriter(f4, "UTF-8")) {
        writer2.println("// AUTO-GENERATED FILE!!!!");
        writer2.println("// This function has yet to be decompiled using 'Dump Current Function' in ghidra");
        writer2.println("//  with possible manualy fixes");
        writer2.println();
        writer2.println("#include <r3/binders/auto.h>");
        writer2.println("#include <r3/binders/stub.h>");
        writer2.println("#include <gh_global.h>");
        writer2.println();
        writer2.println("// " + externalFunction.getEntryPoint());
        writer2.println("// " + externalFunction.getName());
        
        // Parse function signature to extract calling convention, return type, and parameters
        String signature = externalFunction.getSignature().getPrototypeString(false);
        String callingConvention = externalFunction.getCallingConventionName();
        String returnType = externalFunction.getReturnType().toString();
        
        // Generate function stub using appropriate forwarding function
        writer2.println("extern \"C\" " + signature + " {");
        
        // Determine which stub function to use based on calling convention
        String stubFunction;
        if (callingConvention != null && callingConvention.equals("__stdcall")) {
          stubFunction = "gh_stub_impl_stdcall";
        } else {
          // Default to cdecl for most cases
          stubFunction = "gh_stub_impl_cdecl";
        }
        
        // Generate parameter list for the call
        StringBuilder paramList = new StringBuilder();
        var params = externalFunction.getParameters();
        for (int i = 0; i < params.length; i++) {
          if (i > 0) paramList.append(", ");
          paramList.append(params[i].getName());
        }
        
        // Generate the stub call
        if (returnType.equals("void")) {
          writer2.println("  " + stubFunction + "<void>((void*)0x" + 
                         externalFunction.getEntryPoint().toString().replace("0x", "") + 
                         (paramList.length() > 0 ? ", " + paramList.toString() : "") + ");");
        } else {
          writer2.println("  return " + stubFunction + "<" + returnType + ">((void*)0x" + 
                         externalFunction.getEntryPoint().toString().replace("0x", "") + 
                         (paramList.length() > 0 ? ", " + paramList.toString() : "") + ");");
        }
        
        writer2.println("}");
      }

      if (!f4.exists()) {
        createdFile = true;
      }
    }
  }
}