using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Text;
using WhiteMagic.Native;

namespace WhiteMagic.Discovery;

/// <summary>
/// Represents a section in a PE file.
/// </summary>
public readonly record struct PeSection
{
    /// <summary>
    /// The 8-byte null-terminated section name (e.g., ".text", ".data").
    /// </summary>
    public string Name { get; init; }

    /// <summary>
    /// The virtual address of the section when loaded into memory (RVA).
    /// </summary>
    public IntPtr VirtualAddress { get; init; }

    /// <summary>
    /// The size of the section in memory.
    /// </summary>
    public int VirtualSize { get; init; }
}

/// <summary>
/// Parses PE headers to expose section information and entry points.
/// </summary>
public sealed class PeHeaderParser
{
    private readonly MemoryBase _memory;
    private readonly IntPtr _baseAddress;

    /// <summary>
    /// Creates a new PE header parser for the module at the specified base address.
    /// </summary>
    /// <param name="memory">The memory accessor.</param>
    /// <param name="baseAddress">The base address of the module.</param>
    public PeHeaderParser(MemoryBase memory, IntPtr baseAddress)
    {
        ArgumentNullException.ThrowIfNull(memory);
        if (baseAddress == IntPtr.Zero)
            throw new ArgumentException("Base address cannot be zero.", nameof(baseAddress));

        _memory = memory;
        _baseAddress = baseAddress;
    }

    /// <summary>
    /// Gets the entry point RVA (Relative Virtual Address) of the PE file.
    /// </summary>
    /// <returns>The entry point RVA, or <see cref="IntPtr.Zero"/> if unavailable.</returns>
    /// <exception cref="Win32Exception">Reading memory fails.</exception>
    /// <exception cref="InvalidDataException">The PE headers are invalid.</exception>
    public IntPtr EntryPoint
    {
        get
        {
            // Read and parse PE headers
            var (optionalHeader, _) = ParseOptionalHeader();

            if (optionalHeader is null)
                return IntPtr.Zero;

            // Entry point RVA is at offset 16 in the optional header (both PE32 and PE32+)
            return (IntPtr)BitConverter.ToUInt32(optionalHeader.AsSpan(16, 4));
        }
    }

    /// <summary>
    /// Enumerates all sections in the PE file.
    /// </summary>
    /// <returns>An enumerable of PE sections.</returns>
    /// <exception cref="Win32Exception">Reading memory fails.</exception>
    /// <exception cref="InvalidDataException">The PE headers are invalid.</exception>
    public IEnumerable<PeSection> Sections
    {
        get
        {
            var (optionalHeader, sectionHeaders) = ParseOptionalHeaderAndSectionHeaders();

            if (sectionHeaders is null || sectionHeaders.Length == 0)
                yield break;

            foreach (var sectionHeader in sectionHeaders)
            {
                // Parse section name (8-byte, null-terminated)
                string name = ParseSectionName(sectionHeader);

                // VirtualAddress and VirtualSize
                uint virtualAddress = BitConverter.ToUInt32(sectionHeader, 12);
                uint virtualSize = BitConverter.ToUInt32(sectionHeader, 8);

                yield return new PeSection
                {
                    Name = name,
                    VirtualAddress = (IntPtr)virtualAddress,
                    VirtualSize = (int)virtualSize
                };
            }
        }
    }

    /// <summary>
    /// Resolves an exported function's absolute address by name, following export
    /// forwarders (e.g. <c>kernel32!HeapAlloc</c> → <c>NTDLL.RtlAllocateHeap</c>) into
    /// other modules loaded in the same target process.
    /// </summary>
    /// <param name="functionName">The exported symbol name (case-sensitive, as stored
    /// in the export name table).</param>
    /// <returns>The absolute address of the export in the target process.</returns>
    /// <exception cref="InvalidOperationException">The export is not present.</exception>
    /// <exception cref="NotSupportedException">The export forwards to an ordinal or to a
    /// module (such as an API set) that is not resolvable from the target's module list.</exception>
    /// <exception cref="InvalidDataException">The PE export data is malformed.</exception>
    public IntPtr GetExportAddress(string functionName)
    {
        ArgumentException.ThrowIfNullOrEmpty(functionName);
        return ResolveExport(functionName, 0);
    }

    // Maximum forwarder hops before giving up, to bound pathological chains.
    private const int MaxForwarderDepth = 16;

    private IntPtr ResolveExport(string functionName, int depth)
    {
        if (depth > MaxForwarderDepth)
            throw new InvalidDataException($"Export forwarder chain for '{functionName}' is too deep.");

        var (optionalHeader, _) = ParseOptionalHeader();
        if (optionalHeader is null || optionalHeader.Length < 2)
            throw new InvalidDataException("Optional header unavailable.");

        // 0x10b = PE32 (32-bit), 0x20b = PE32+ (64-bit). Data directories start at a
        // different offset in each: 96 for PE32, 112 for PE32+. The export table is
        // directory index 0, so its 8-byte entry sits at that offset.
        ushort magic = BitConverter.ToUInt16(optionalHeader, 0);
        bool pe32Plus = magic == 0x20b;
        int exportDirOffset = pe32Plus ? 112 : 96;
        if (optionalHeader.Length < exportDirOffset + 8)
            throw new InvalidDataException("Optional header does not contain the export data directory.");

        uint exportRva = BitConverter.ToUInt32(optionalHeader, exportDirOffset);
        uint exportSize = BitConverter.ToUInt32(optionalHeader, exportDirOffset + 4);
        if (exportRva == 0 || exportSize == 0)
            throw new InvalidOperationException("Module has no export table.");

        // IMAGE_EXPORT_DIRECTORY is 40 bytes.
        byte[] dir = _memory.ReadBytes(_baseAddress + (nint)exportRva, 40);
        if (dir.Length < 40)
            throw new InvalidDataException("Failed to read the export directory.");

        uint numberOfFunctions = BitConverter.ToUInt32(dir, 20);
        uint numberOfNames = BitConverter.ToUInt32(dir, 24);
        uint addressOfFunctions = BitConverter.ToUInt32(dir, 28);
        uint addressOfNames = BitConverter.ToUInt32(dir, 32);
        uint addressOfNameOrdinals = BitConverter.ToUInt32(dir, 36);

        // Guard against corrupt counts before allocating arrays sized from them.
        if (numberOfNames > 0x10000 || numberOfFunctions > 0x10000)
            throw new InvalidDataException("Export table entry count is out of range.");

        if (numberOfNames == 0)
            throw new InvalidOperationException($"Export '{functionName}' not found (module exports no names).");

        byte[] nameRvas = _memory.ReadBytes(_baseAddress + (nint)addressOfNames, checked((int)(numberOfNames * 4)));
        byte[] nameOrdinals = _memory.ReadBytes(_baseAddress + (nint)addressOfNameOrdinals, checked((int)(numberOfNames * 2)));
        if (nameRvas.Length < numberOfNames * 4 || nameOrdinals.Length < numberOfNames * 2)
            throw new InvalidDataException("Failed to read the export name tables.");

        int nameIndex = -1;
        for (int i = 0; i < numberOfNames; i++)
        {
            uint nameRva = BitConverter.ToUInt32(nameRvas, i * 4);
            string name = _memory.ReadString(_baseAddress + (nint)nameRva, Encoding.ASCII, 512);
            if (string.Equals(name, functionName, StringComparison.Ordinal))
            {
                nameIndex = i;
                break;
            }
        }

        if (nameIndex < 0)
            throw new InvalidOperationException($"Export '{functionName}' not found in module.");

        ushort ordinal = BitConverter.ToUInt16(nameOrdinals, nameIndex * 2);
        if (ordinal >= numberOfFunctions)
            throw new InvalidDataException("Export name ordinal is out of range.");

        byte[] funcRvaBytes = _memory.ReadBytes(
            _baseAddress + (nint)(addressOfFunctions + (uint)ordinal * 4u), 4);
        if (funcRvaBytes.Length < 4)
            throw new InvalidDataException("Failed to read the export address table entry.");

        uint funcRva = BitConverter.ToUInt32(funcRvaBytes, 0);
        if (funcRva == 0)
            throw new InvalidOperationException($"Export '{functionName}' has no address.");

        // A function RVA that lands inside the export directory region is not code but a
        // null-terminated "Module.Function" forwarder string.
        if (funcRva >= exportRva && funcRva < exportRva + exportSize)
        {
            string forwarder = _memory.ReadString(_baseAddress + (nint)funcRva, Encoding.ASCII, 512);
            return ResolveForwarder(forwarder, depth);
        }

        return _baseAddress + (nint)funcRva;
    }

    private IntPtr ResolveForwarder(string forwarder, int depth)
    {
        int dot = forwarder.LastIndexOf('.');
        if (dot <= 0 || dot >= forwarder.Length - 1)
            throw new InvalidDataException($"Malformed export forwarder string '{forwarder}'.");

        string moduleName = forwarder[..dot];
        string target = forwarder[(dot + 1)..];

        if (target.StartsWith('#'))
        {
            throw new NotSupportedException(
                $"Ordinal export forwarders are not supported (forwarder '{forwarder}').");
        }

        IntPtr targetBase = RemoteModule.ResolveBase(_memory.ProcessId, moduleName);
        if (targetBase == IntPtr.Zero)
        {
            throw new NotSupportedException(
                $"Export forwarder target module '{moduleName}' is not loaded in the target " +
                $"process, or is an unresolvable API set (forwarder '{forwarder}').");
        }

        return new PeHeaderParser(_memory, targetBase).ResolveExport(target, depth + 1);
    }

    /// <summary>
    /// Parses the DOS header, PE signature, and optional header.
    /// </summary>
    private (byte[]? OptionalHeader, byte[][]? SectionHeaders) ParseOptionalHeaderAndSectionHeaders()
    {
        // Read DOS header (first 64 bytes)
        byte[] dosHeader = _memory.ReadBytes(_baseAddress, 64);
        if (dosHeader.Length < 64)
            throw new InvalidDataException("Failed to read DOS header.");

        // Verify DOS signature "MZ"
        if (dosHeader[0] != 0x4D || dosHeader[1] != 0x5A)
            throw new InvalidDataException("Invalid DOS signature (not a PE file).");

        // PE header offset is at 0x3C in DOS header
        int peOffset = BitConverter.ToInt32(dosHeader, 0x3C);
        if (peOffset < 0 || peOffset > 0x1000) // Sanity check
            throw new InvalidDataException($"Invalid PE offset: {peOffset}");

        // Read PE signature (4 bytes: "PE\0\0")
        IntPtr peSigAddr = _baseAddress + peOffset;
        byte[] peSignature = _memory.ReadBytes(peSigAddr, 4);
        if (peSignature.Length < 4)
            throw new InvalidDataException("Failed to read PE signature.");

        if (peSignature[0] != 0x50 || peSignature[1] != 0x45 ||
            peSignature[2] != 0x00 || peSignature[3] != 0x00)
            throw new InvalidDataException("Invalid PE signature.");

        // COFF header follows PE signature (20 bytes)
        IntPtr coffAddr = peSigAddr + 4;
        byte[] coffHeader = _memory.ReadBytes(coffAddr, 20);
        if (coffHeader.Length < 20)
            throw new InvalidDataException("Failed to read COFF header.");

        // SizeOfOptionalHeader is at offset 16 in COFF header
        ushort sizeOfOptionalHeader = BitConverter.ToUInt16(coffHeader, 16);
        // NumberOfSections is at offset 2 in COFF header
        ushort numberOfSections = BitConverter.ToUInt16(coffHeader, 2);

        if (numberOfSections == 0 || numberOfSections > 96)
            return (null, null); // No sections or unreasonable number
        // Optional header follows COFF header
        IntPtr optAddr = coffAddr + 20;
        byte[] optionalHeader = _memory.ReadBytes(optAddr, sizeOfOptionalHeader);
        if (optionalHeader.Length < sizeOfOptionalHeader)
            throw new InvalidDataException("Failed to read optional header.");

        // Section headers follow optional header
        IntPtr sectionAddr = optAddr + sizeOfOptionalHeader;
        int sectionHeaderSize = 40; // IMAGE_SECTION_HEADER is 40 bytes

        byte[][] sectionHeaders = new byte[numberOfSections][];
        for (int i = 0; i < numberOfSections; i++)
        {
            byte[] section = _memory.ReadBytes(sectionAddr + (i * sectionHeaderSize), sectionHeaderSize);
            if (section.Length < sectionHeaderSize)
                throw new InvalidDataException($"Failed to read section header {i}.");
            sectionHeaders[i] = section;
        }

        return (optionalHeader, sectionHeaders);
    }

    /// <summary>
    /// Parses just the optional header (for entry point).
    /// </summary>
    private (byte[]? OptionalHeader, byte[][]? SectionHeaders) ParseOptionalHeader()
    {
        // Read DOS header (first 64 bytes)
        byte[] dosHeader = _memory.ReadBytes(_baseAddress, 64);
        if (dosHeader.Length < 64)
            throw new InvalidDataException("Failed to read DOS header.");

        // Verify DOS signature "MZ"
        if (dosHeader[0] != 0x4D || dosHeader[1] != 0x5A)
            throw new InvalidDataException("Invalid DOS signature (not a PE file).");

        // PE header offset is at 0x3C in DOS header
        int peOffset = BitConverter.ToInt32(dosHeader, 0x3C);
        if (peOffset < 0 || peOffset > 0x1000) // Sanity check
            throw new InvalidDataException($"Invalid PE offset: {peOffset}");

        // Read PE signature (4 bytes: "PE\0\0")
        IntPtr peSigAddr = _baseAddress + peOffset;
        byte[] peSignature = _memory.ReadBytes(peSigAddr, 4);
        if (peSignature.Length < 4)
            throw new InvalidDataException("Failed to read PE signature.");

        if (peSignature[0] != 0x50 || peSignature[1] != 0x45 ||
            peSignature[2] != 0x00 || peSignature[3] != 0x00)
            throw new InvalidDataException("Invalid PE signature.");

        // COFF header follows PE signature (20 bytes)
        IntPtr coffAddr = peSigAddr + 4;
        byte[] coffHeader = _memory.ReadBytes(coffAddr, 20);
        if (coffHeader.Length < 20)
            throw new InvalidDataException("Failed to read COFF header.");

        // SizeOfOptionalHeader is at offset 16 in COFF header
        ushort sizeOfOptionalHeader = BitConverter.ToUInt16(coffHeader, 16);

        // Optional header follows COFF header
        IntPtr optAddr = coffAddr + 20;
        byte[] optionalHeader = _memory.ReadBytes(optAddr, sizeOfOptionalHeader);
        if (optionalHeader.Length < sizeOfOptionalHeader)
            throw new InvalidDataException("Failed to read optional header.");

        return (optionalHeader, null);
    }

    /// <summary>
    /// Determines whether the PE file is PE32+ (64-bit) or PE32 (32-bit).
    /// </summary>
    private bool IsPe32Plus()
    {
        var (optionalHeader, _) = ParseOptionalHeaderAndSectionHeaders();

        if (optionalHeader is null || optionalHeader.Length < 2)
            throw new InvalidDataException("Optional header too short.");

        // Magic is at offset 0 in optional header
        // 0x10b = PE32 (32-bit), 0x20b = PE32+ (64-bit)
        ushort magic = BitConverter.ToUInt16(optionalHeader, 0);
        return magic == 0x20b;
    }

    /// <summary>
    /// Parses a null-terminated 8-byte section name.
    /// </summary>
    private static string ParseSectionName(byte[] sectionHeader)
    {
        // Name is first 8 bytes
        var nameBytes = new Span<byte>(sectionHeader, 0, 8);

        // Find null terminator
        int len = 0;
        for (; len < 8; len++)
        {
            if (nameBytes[len] == 0)
                break;
        }

        return System.Text.Encoding.ASCII.GetString(nameBytes[..len]);
    }
}