using X86Disassembler.PE;
using X86Disassembler.X86;

namespace X86Disassembler;

/// <summary>
/// Main program class
/// </summary>
public class Program
{
    // Hardcoded file path for testing
    private const string FilePath = @"C:\Program Files (x86)\Nikita\Iron Strategy\Terrain.dll";
    
    /// <summary>
    /// Main entry point
    /// </summary>
    /// <param name="args">Command line arguments</param>
    public static void Main(string[] args)
    {
        Console.WriteLine("X86 Disassembler and Decompiler");
        Console.WriteLine("--------------------------------");
        
        // Load the file
        Console.WriteLine($"Loading file: {FilePath}");
        byte[] fileBytes = File.ReadAllBytes(FilePath);
        Console.WriteLine($"Successfully loaded {FilePath}");
        Console.WriteLine($"File size: {fileBytes.Length} bytes\n");
        
        // Parse the PE format
        Console.WriteLine("Parsing PE format...\n");
        PeFile peFile = new PeFile(fileBytes);
        peFile.Parse();
        
        // Print PE file information
        Console.WriteLine("PE File Information:");
        Console.WriteLine($"Architecture: {(peFile.OptionalHeader.Is64Bit() ? "64-bit" : "32-bit")}");
        Console.WriteLine($"Entry Point: 0x{peFile.OptionalHeader.AddressOfEntryPoint:X8}");
        Console.WriteLine($"Image Base: 0x{peFile.OptionalHeader.ImageBase:X8}");
        Console.WriteLine($"Number of Sections: {peFile.FileHeader.NumberOfSections}");
        Console.WriteLine();
        
        // Print section information
        PrintPeSections(peFile);

        // Print export information
        PrintPeExports(peFile);

        // Print import information
        PrintPeImports(peFile);
        
        // Find code sections
        var codeSections = peFile.SectionHeaders.FindAll(s => s.ContainsCode());
        Console.WriteLine($"Found {codeSections.Count} code section(s):");
        foreach (var section in codeSections)
        {
            Console.WriteLine($"  - {section.Name}: Size={section.VirtualSize} bytes, RVA=0x{section.VirtualAddress:X8}");
        }
        Console.WriteLine();
        
        // Disassemble the first code section
        if (codeSections.Count > 0)
        {
            var section = codeSections[0];
            byte[] codeBytes = peFile.GetSectionData(peFile.SectionHeaders.IndexOf(section));
            
            // First demonstrate sequential disassembly
            Console.WriteLine($"Sequential disassembly of section {section.Name} at RVA 0x{section.VirtualAddress:X8}:");
            
            // Create a disassembler for the code section
            // Base address should be the section's virtual address, not the image base + VA
            Disassembler disassembler = new Disassembler(codeBytes, section.VirtualAddress);
            
            // Disassemble sequentially (linear approach)
            var linearInstructions = disassembler.Disassemble();
            
            // Print the first 30 instructions from linear disassembly
            int linearCount = Math.Min(30, linearInstructions.Count);
            for (int i = 0; i < linearCount; i++)
            {
                Console.WriteLine(linearInstructions[i]);
            }
            
            // Print a summary of how many more instructions there are
            if (linearInstructions.Count > linearCount)
            {
                Console.WriteLine($"... ({linearInstructions.Count - linearCount} more instructions not shown)");
            }
            
            Console.WriteLine();
            Console.WriteLine("====================================================");
            Console.WriteLine();
            
            // Now demonstrate control flow-based disassembly from entry point
            Console.WriteLine($"Control flow-based disassembly starting from entry point 0x{peFile.OptionalHeader.AddressOfEntryPoint:X8}:");
            
            try
            {
                // Get the entry point RVA from the PE header
                uint entryPointRva = peFile.OptionalHeader.AddressOfEntryPoint;
                
                // Make sure the entry point is within this code section
                if (entryPointRva >= section.VirtualAddress && 
                    entryPointRva < section.VirtualAddress + section.VirtualSize)
                {
                    // Disassemble starting from the entry point (control flow-based)
                    var cfgInstructions = disassembler.DisassembleFunction(entryPointRva);
                    
                    // Print the instructions from the entry point function
                    int cfgCount = Math.Min(50, cfgInstructions.Count);
                    for (int i = 0; i < cfgCount; i++)
                    {
                        Console.WriteLine(cfgInstructions[i]);
                    }
                    
                    // Print a summary if there are more instructions
                    if (cfgInstructions.Count > cfgCount)
                    {
                        Console.WriteLine($"... ({cfgInstructions.Count - cfgCount} more instructions in this function not shown)");
                    }
                    
                    Console.WriteLine();
                    Console.WriteLine($"Found {cfgInstructions.Count} instructions following control flow from entry point.");
                }
                else
                {
                    // Try one of the exported functions instead
                    Console.WriteLine($"Entry point is not in the {section.Name} section. Trying the first exported function instead...");
                    
                    if (peFile.ExportDirectory != null && peFile.ExportedFunctions.Count > 0)
                    {
                        uint functionRva = peFile.ExportedFunctions[0].AddressRva;
                        Console.WriteLine($"Disassembling exported function at RVA 0x{functionRva:X8} ({peFile.ExportedFunctions[0].Name}):");
                        
                        var cfgInstructions = disassembler.DisassembleFunction(functionRva);
                        
                        // Print the instructions from the function
                        int cfgCount = Math.Min(50, cfgInstructions.Count);
                        for (int i = 0; i < cfgCount; i++)
                        {
                            Console.WriteLine(cfgInstructions[i]);
                        }
                        
                        // Print a summary if there are more instructions
                        if (cfgInstructions.Count > cfgCount)
                        {
                            Console.WriteLine($"... ({cfgInstructions.Count - cfgCount} more instructions in this function not shown)");
                        }
                        
                        Console.WriteLine();
                        Console.WriteLine($"Found {cfgInstructions.Count} instructions following control flow from exported function.");
                    }
                    else
                    {
                        Console.WriteLine("No exported functions found to disassemble.");
                    }
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Error during control flow disassembly: {ex.Message}");
            }
        }
        
        // Console.WriteLine("\nPress Enter to exit...");
        // Console.ReadLine();
    }

    private static void PrintPeImports(PeFile peFile)
    {
        Console.WriteLine("Imported Functions:");
        Console.WriteLine($"Number of Imported DLLs: {peFile.ImportDescriptors.Count}");

        foreach (var import in peFile.ImportDescriptors)
        {
            Console.WriteLine($"  DLL: {import.DllName}");
                
            for (int i = 0; i < import.Functions.Count; i++)
            {
                var function = import.Functions[i];
                if (function.IsOrdinal)
                {
                    Console.WriteLine($"    {i}: Ordinal {function.Ordinal}");
                }
                else
                {
                    Console.WriteLine($"    {i}: {function.Name} (Hint={function.Hint})");
                }
            }
        }
        Console.WriteLine();
    }

    private static void PrintPeExports(PeFile peFile)
    {
        Console.WriteLine("Exported Functions:");
        Console.WriteLine($"DLL Name: {peFile.ExportDirectory.DllName}");
        Console.WriteLine($"Number of Functions: {peFile.ExportDirectory.NumberOfFunctions}");
        Console.WriteLine($"Number of Names: {peFile.ExportDirectory.NumberOfNames}");
            
        for (int i = 0; i < peFile.ExportedFunctions.Count; i++)
        {
            var export = peFile.ExportedFunctions[i];
            Console.WriteLine($"  {i}: {export.Name} (Ordinal={export.Ordinal}, RVA=0x{export.AddressRva:X8})");
        }
        Console.WriteLine();
    }

    private static void PrintPeSections(PeFile peFile)
    {
        Console.WriteLine("Sections:");
        foreach (var section in peFile.SectionHeaders)
        {
            string flags = "";
            if (section.ContainsCode()) flags += "Code ";
            if (section.IsExecutable()) flags += "Exec ";
            if (section.IsReadable()) flags += "Read ";
            if (section.IsWritable()) flags += "Write";
            
            Console.WriteLine($"  {peFile.SectionHeaders.IndexOf(section)}: {section.Name,-8} VA=0x{section.VirtualAddress:X8} Size={section.VirtualSize,-8} [{flags}]");
        }
        Console.WriteLine();
    }
}