// Source code is decompiled from a .class file using FernFlower decompiler.
package re3lib;

import ghidra.program.model.data.*;
import ghidra.program.model.data.Enum;
import ghidra.util.Msg;
import ghidra.util.exception.CancelledException;
import ghidra.util.task.TaskMonitor;
import java.io.IOException;
import java.io.Writer;
import java.util.*;

public class DataTypeWriter {
  public HashSet<String> requiredFunctionTypes;
  public HashMap<String, FunctionDefinition> lazyFunctionTypeBlocks;

  private static String[] INTEGRAL_TYPES = new String[] { "char", "short", "int", "long", "long long", "__int64",
      "float", "double", "long double", "void" };
  private static String[] INTEGRAL_MODIFIERS = new String[] { "signed", "unsigned", "const", "static", "volatile",
      "mutable" };
  private static String EOL = System.getProperty("line.separator");

  private Writer writer;
  private DataTypeManager dtm;
  private DataOrganization dataOrganization;
  private AnnotationHandler annotator;
  private boolean cppStyleComments;

  // Block represents a single data type declaration with its dependencies
  private static class Block {
    final DataType dataType;
    final String code;
    final Set<String> dependencies;

    Block(DataType dataType, String code, Set<String> dependencies) {
      this.dataType = dataType;
      this.code = code;
      this.dependencies = new HashSet<>(dependencies);
    }
  }

  public DataTypeWriter(DataTypeManager dtm, Writer writer) throws IOException {
    this(dtm, writer, new DefaultAnnotationHandler());
  }

  public DataTypeWriter(DataTypeManager dtm, Writer writer, boolean cppStyleComments) throws IOException {
    this(dtm, writer, new DefaultAnnotationHandler(), cppStyleComments);
  }

  public DataTypeWriter(DataTypeManager dtm, Writer writer, AnnotationHandler annotator) throws IOException {
    this(dtm, writer, annotator, false);
  }

  public DataTypeWriter(DataTypeManager dtm, Writer writer, AnnotationHandler annotator, boolean cppStyleComments)
      throws IOException {
    this.cppStyleComments = false;
    this.dtm = dtm;
    if (dtm != null) {
      this.dataOrganization = dtm.getDataOrganization();
    }

    if (this.dataOrganization == null) {
      this.dataOrganization = DataOrganizationImpl.getDefaultOrganization();
    }

    this.writer = writer;
    this.annotator = annotator;
    this.cppStyleComments = cppStyleComments;
    this.requiredFunctionTypes = new HashSet<>();
    this.lazyFunctionTypeBlocks = new HashMap<>();
  }

  private String comment(String text) {
    if (text == null) {
      return "";
    } else {
      return this.cppStyleComments ? "// " + text : "/* " + text + " */";
    }
  }

  public void write(DataTypeManager dataTypeManager, TaskMonitor monitor) throws IOException, CancelledException {
    this.write(dataTypeManager.getRootCategory(), monitor);
  }

  public void write(Category category, TaskMonitor monitor) throws IOException, CancelledException {
    DataType[] dataTypes = category.getDataTypes();
    this.write(dataTypes, monitor);
    Category[] subCategories = category.getCategories();
    Category[] var5 = subCategories;
    int var6 = subCategories.length;

    for (int var7 = 0; var7 < var6; ++var7) {
      Category subCategory = var5[var7];
      if (monitor.isCancelled()) {
        return;
      }

      this.write(subCategory, monitor);
    }
  }

  public void write(DataType[] dataTypes, TaskMonitor monitor) throws IOException, CancelledException {
    write(Arrays.asList(dataTypes), monitor);
  }

  public void write(List<DataType> dataTypes, TaskMonitor monitor) throws IOException, CancelledException {
    this.write(dataTypes, monitor, true);
  }

  public void write(List<DataType> dataTypes, TaskMonitor monitor, boolean throwExceptionOnInvalidType)
      throws IOException, CancelledException {
    monitor.initialize((long) dataTypes.size());

    // Step 1: Create blocks for each data type
    Map<String, Block> blocks = new HashMap<>();

    int cnt = 0;
    for (DataType dataType : dataTypes) {
      monitor.checkCancelled();

      if (dataType == null) {
        continue;
      }

      try {
        Block block = createBlock(dataType);
        if (block != null) {
          blocks.put(dataType.getDisplayName(), block);
        }
      } catch (Exception e) {
        if (throwExceptionOnInvalidType) {
          throw new IOException("Failed to process data type: " + dataType.getDisplayName(), e);
        }
        Msg.error(this, "Failed to process data type: " + dataType.getDisplayName(), e);
      }

      ++cnt;
      monitor.setProgress((long) cnt);
    }

    // Step 2: Check if we're missing any function type typedefs
    for (String funcName : requiredFunctionTypes) {
      if (!blocks.containsKey(funcName)) {
        FunctionDefinition funcType = lazyFunctionTypeBlocks.get(funcName);

        Set<String> dependencies = new HashSet<>();
        StringBuilder code = new StringBuilder();
        this.writeFunctionDefinitionBlock(funcType, funcName, code, dependencies);
        blocks.put(funcName, new Block(funcType, code.toString(), dependencies));
      }
    }
    

    // Step 3: Topological sort and write
    List<Block> sortedBlocks = topologicalSort(blocks);
    for (Block block : sortedBlocks) {
      writer.write(block.code);
      writer.write(EOL);
    }

    writer.flush();
  }

  private Block createBlock(DataType dt) throws IOException {
    if (dt instanceof FactoryDataType) {
      return null; // Skip only factory types
    }

    dt = dt.clone(this.dtm);
    Set<String> dependencies = new HashSet<>();
    StringBuilder code = new StringBuilder();

    // if (dt.getDisplayName().contains("HIE_tduLinkedObject")) {
    //   System.out.println("DEBUG " + dt.getDisplayName());
    // }

    if (dt.equals(DataType.DEFAULT)) {
      code.append("typedef unsigned char   ").append(DataType.DEFAULT.getName()).append(";");
    } else if (dt instanceof Dynamic) {
      // Handle dynamic types
      Dynamic dynamicType = (Dynamic) dt;
      DataType baseDt = dynamicType.getReplacementBaseType();
      if (baseDt != null) {
        dependencies.add(baseDt.getDisplayName());
      }
      code.append(comment("Dynamic type: " + dt.getDisplayName()));
    } else if (dt instanceof Structure) {
      writeStructureBlock((Structure) dt, code, dependencies);
    } else if (dt instanceof Union) {
      writeUnionBlock((Union) dt, code, dependencies);
    } else if (dt instanceof Enum) {
      if(dt.getDisplayName().contains(".conflict"))
        return null;
      writeEnumBlock((Enum) dt, code, dependencies);
    } else if (dt instanceof TypeDef) {
      writeTypeDefBlock((TypeDef) dt, code, dependencies);
    } else if (dt instanceof BuiltInDataType) {
      writeBuiltInBlock((BuiltInDataType) dt, code, dependencies);
    } else if (dt instanceof FunctionDefinition) {
      // Store these for later, a lot of them are going to be unused
      lazyFunctionTypeBlocks.put(dt.getDisplayName(), (FunctionDefinition) dt);
      return null;
    } else {
      code.append(comment("Unable to write datatype. Type unrecognized: " + dt.getClass()));
    }

    return new Block(dt, code.toString(), dependencies);
  }

  private void writeStructureBlock(Structure struct, StringBuilder code, Set<String> dependencies) {
    String structName = struct.getDisplayName();

    // Struct definition
    code.append("struct ").append(structName).append(" {");
    String descrip = struct.getDescription();
    if (descrip != null && descrip.length() > 0) {
      code.append(" ").append(comment(descrip));
    }
    code.append(EOL);

    // Process components
    for (DataTypeComponent component : struct.getComponents()) {
      writeComponentBlock(component, struct, code, dependencies);
    }

    code.append(annotator.getSuffix(struct, null));
    code.append("};");
  }

  private void writeUnionBlock(Union union, StringBuilder code, Set<String> dependencies) {
    String unionName = union.getDisplayName();

    // Union definition (no forward declaration needed - let dependency ordering
    // handle it)
    code.append("union ").append(unionName).append(" {");
    String descrip = union.getDescription();
    if (descrip != null && descrip.length() > 0) {
      code.append(" ").append(comment(descrip));
    }
    code.append(EOL);

    // Process components
    for (DataTypeComponent component : union.getComponents()) {
      writeComponentBlock(component, union, code, dependencies);
    }

    code.append(annotator.getSuffix(union, null));
    code.append("};");
  }

  private void writeComponentBlock(DataTypeComponent component, Composite composite, StringBuilder code,
      Set<String> dependencies) {
    code.append("    ");
    code.append(annotator.getPrefix(composite, component));

    String fieldName = component.getFieldName();
    String originalName = fieldName;
    boolean needsFixing = fieldName != null && fieldName.contains("?");
    if (fieldName == null || fieldName.length() == 0 || needsFixing) {
      fieldName = component.getDefaultFieldName();
    }

    DataType componentDataType = component.getDataType();

    DataType depType = getImmediateDependencyType(componentDataType);
    if (depType != null) {
      dependencies.add(depType.getDisplayName());
      if(depType instanceof FunctionDefinition) {
        requiredFunctionTypes.add(depType.getDisplayName());
      }
    }

    code.append(getTypeDeclaration(fieldName, componentDataType, component.getLength()));
    code.append(";");
    code.append(annotator.getSuffix(composite, component));

    String comment = component.getComment();
    String commentText = comment != null && comment.length() > 0 ? comment : "";
    if (needsFixing) {
      commentText += (commentText.length() > 0 ? " " : "") + "(fix name \"" + originalName + "\")";
    }
    if (commentText.length() > 0) {
      code.append(" ").append(comment(commentText));
    }
    code.append(EOL);
  }

  private void writeEnumBlock(Enum enumm, StringBuilder code, Set<String> dependencies) {
    String enumName = enumm.getDisplayName();

    if (enumName.startsWith("define_") && enumName.length() > 7 && enumm.getCount() == 1) {
      long val = enumm.getValues()[0];
      code.append("#define ").append(enumName.substring(7)).append(" ").append(Long.toString(val));
    } else {
      code.append("enum ").append(enumName).append(" {");
      String description = enumm.getDescription();
      if (description != null && description.length() != 0) {
        code.append(" ").append(comment(description));
      }
      code.append(EOL);

      String[] names = enumm.getNames();
      for (int j = 0; j < names.length; ++j) {
        code.append("    ");
        code.append(annotator.getPrefix(enumm, names[j]));
        code.append(names[j]);
        code.append("=");
        code.append(Long.toString(enumm.getValue(names[j])));
        String comment = enumm.getComment(names[j]);
        if (!comment.isBlank()) {
          code.append(" ").append(comment(comment));
        }
        code.append(annotator.getSuffix(enumm, names[j]));
        if (j < names.length - 1) {
          code.append(",");
        }
        code.append(EOL);
      }
      code.append("};");
    }
  }

  private boolean writeFunctionDefinitionBlock(FunctionDefinition funcDef, String name, StringBuilder code,
      Set<String> dependencies) {
    DataType returnType = funcDef.getReturnType();
    ParameterDefinition[] params = funcDef.getArguments();

    // Add return type dependency if not built-in
    if (returnType != null && !isBuiltInType(returnType)) {
      DataType depType = getImmediateDependencyType(returnType);
      if (depType != null) {
        dependencies.add(depType.getDisplayName());
      }
    }

    // Build function typedef
    code.append("typedef ");
    if (returnType != null) {
      code.append(getTypeDeclaration("", returnType, -1)).append(" ");
    } else {
      code.append("void ");
    }

    code.append("(*").append(name).append(")(");

    if (params != null && params.length > 0) {
      for (int i = 0; i < params.length; i++) {
        ParameterDefinition param = params[i];
        DataType paramType = param.getDataType();

        // Add parameter type dependency if not built-in
        if (!isBuiltInType(paramType) && !isPointerType(paramType)) {
          DataType depType = getImmediateDependencyType(paramType);
          if (depType != null) {
            dependencies.add(depType.getDisplayName());
          }
        }

        String paramName = param.getName();
        if (paramName == null || paramName.isEmpty()) {
          paramName = "param" + (i + 1);
        }

        code.append(getTypeDeclaration(paramName, paramType, -1));
        if (i < params.length - 1) {
          code.append(", ");
        }
      }
    } else {
      code.append("void");
    }

    code.append(");");
    return true;
  }

  private boolean writeTypeDefBlock(TypeDef typeDef, StringBuilder code, Set<String> dependencies) {
    String typedefName = typeDef.getDisplayName();
    DataType dataType = typeDef.getDataType();

    // Handle function definition typedefs
    if (dataType instanceof FunctionDefinition) {
      return writeFunctionDefinitionBlock((FunctionDefinition) dataType, typedefName, code, dependencies);
    } // Could be pointer to function
    else if (dataType instanceof Pointer) {
      Pointer ptr = (Pointer) dataType;
      if (ptr.getDataType() instanceof FunctionDefinition) {
        return writeFunctionDefinitionBlock((FunctionDefinition) ((Pointer) dataType).getDataType(), typedefName, code,
            dependencies);
      } else {
        // Try to write the typedef as a forward declaration
        code.append("typedef ");
        code.append(getTypeDeclaration(typedefName, dataType, -1));
        code.append(";");
        return true;
      }
    } else {
      DataType depType = getImmediateDependencyType(dataType);
      if (depType != null && !isBuiltInType(depType)) {
        dependencies.add(depType.getDisplayName());
      }

      String typedefString = getTypeDeclaration(typedefName, dataType, -1);
      code.append("typedef ").append(typedefString).append(";");
      return true;
    }
  }

  private void writeBuiltInBlock(BuiltInDataType dt, StringBuilder code, Set<String> dependencies) {
    String declaration = dt.getCTypeDeclaration(this.dataOrganization);
    if (declaration != null) {
      code.append(declaration);
    }
  }

  private boolean isPointerType(DataType dt) {
    return dt instanceof Pointer;
  }

  private boolean isBuiltInType(DataType dt) {
    return dt instanceof BuiltInDataType;
  }

  private DataType getImmediateDependencyType(DataType dt) {
    // Only unwrap arrays and bitfields, but preserve typedefs and other types
    while (dt != null) {
      if (dt instanceof Array) {
        dt = ((Array) dt).getDataType();
      } else if (dt instanceof BitFieldDataType) {
        dt = ((BitFieldDataType) dt).getBaseDataType();
      } else if (dt instanceof Pointer) {
        dt = ((Pointer) dt).getDataType();
      } else {
        break;
      }
    }
    return dt;
  }

  private String getTypeDeclaration(String name, DataType dataType, int instanceLength) {
    if (name == null) {
      name = "";
    }

    StringBuilder sb = new StringBuilder();

    if (dataType instanceof BitFieldDataType) {
      BitFieldDataType bfDt = (BitFieldDataType) dataType;
      name = name + ":" + bfDt.getDeclaredBitSize();
      dataType = bfDt.getBaseDataType();
    }

    // Handle arrays and pointers
    while (dataType instanceof Array) {
      Array array = (Array) dataType;
      name = name + "[" + array.getNumElements() + "]";
      dataType = array.getDataType();
    }

    while (dataType instanceof Pointer) {
      Pointer pointer = (Pointer) dataType;
      DataType elem = pointer.getDataType();
      if (elem == null) {
        break;
      }
      name = "*" + name;
      dataType = elem;
      if (elem instanceof Array) {
        name = "(" + name + ")";
      }
    }

    String prefix = getDataTypePrefix(dataType);

    String dataTypeString;
    if (dataType instanceof AbstractIntegerDataType) {
      dataTypeString = ((AbstractIntegerDataType) dataType).getCDeclaration();
    } else {
      dataTypeString = dataType.getDisplayName();
    }

    String componentString = prefix + dataTypeString;
    if (name.length() != 0) {
      componentString = componentString + " " + name;
    }

    return componentString;
  }

  private String getDataTypePrefix(DataType dataType) {
    // Don't add struct/union prefix for typedefs - they already have their own name
    if (dataType instanceof TypeDef) {
      return "";
    }

    // Only add struct/union prefix for direct struct/union references (not
    // typedefs)
    if (dataType instanceof Structure) {
      return "struct ";
    } else if (dataType instanceof Union) {
      return "union ";
    } else if (dataType instanceof Enum) {
      return "enum ";
    } else {
      return "";
    }
  }

  private boolean isIntegral(String typedefName, String basetypeName) {
    for (String type : INTEGRAL_TYPES) {
      if (typedefName.equals(type)) {
        return true;
      }
    }

    boolean endsWithIntegralType = false;
    for (String type : INTEGRAL_TYPES) {
      if (typedefName.endsWith(" " + type)) {
        endsWithIntegralType = true;
        break;
      }
    }

    for (String modifier : INTEGRAL_MODIFIERS) {
      if (typedefName.indexOf(modifier + " ") >= 0 || typedefName.indexOf(" " + modifier) >= 0) {
        return true;
      }
    }

    if (endsWithIntegralType) {
      return true;
    } else if (typedefName.endsWith(" " + basetypeName)) {
      return false; // Let it through as a regular typedef
    } else {
      return false;
    }
  }

  private List<Block> topologicalSort(Map<String, Block> blocks) {
    List<Block> result = new ArrayList<>();
    Set<String> visited = new HashSet<>();
    Set<String> visiting = new HashSet<>();

    for (Block block : blocks.values()) {
      if (!visited.contains(block.dataType.getDisplayName())) {
        topologicalSortVisit(block, blocks, visited, visiting, result);
      }
    }

    return result;
  }

  private void topologicalSortVisit(Block block, Map<String, Block> blocks,
      Set<String> visited, Set<String> visiting, List<Block> result) {
    String blockName = block.dataType.getDisplayName();

    if (visiting.contains(blockName)) {
      // Circular dependency detected, but we'll continue (forward declarations should
      // handle this)
      return;
    }

    if (visited.contains(blockName)) {
      return;
    }

    visiting.add(blockName);

    // Visit dependencies first
    for (String dep : block.dependencies) {
      Block depBlock = blocks.get(dep);
      if (depBlock != null) {
        topologicalSortVisit(depBlock, blocks, visited, visiting, result);
      }
    }

    visiting.remove(blockName);
    visited.add(blockName);
    result.add(block);
  }
}
