using System.Text;

namespace ChangeT50.Adapter;

public sealed record PortableExecutableInspection(
    ushort Machine,
    bool IsPe32,
    string ExportedModuleName,
    IReadOnlyList<string> ExportNames);

/// <summary>
/// Reads just enough of a Windows PE file to verify the supplied DLL on non-Windows hosts.
/// It never loads or executes the inspected file.
/// </summary>
public static class PortableExecutableInspector
{
    private const uint PeSignature = 0x00004550;
    private const ushort Pe32Magic = 0x010b;
    private const int MaximumExportNames = 100_000;

    public static PortableExecutableInspection Inspect(string path)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(path);

        using var stream = File.OpenRead(path);
        using var reader = new BinaryReader(stream, Encoding.ASCII, leaveOpen: false);

        if (stream.Length < 64 || reader.ReadUInt16() != 0x5a4d)
        {
            throw new InvalidDataException("File is not an MZ executable.");
        }

        stream.Position = 0x3c;
        var peOffset = reader.ReadUInt32();
        EnsureRange(stream, peOffset, 24, "PE header");
        stream.Position = peOffset;

        if (reader.ReadUInt32() != PeSignature)
        {
            throw new InvalidDataException("PE signature was not found.");
        }

        var machine = reader.ReadUInt16();
        var sectionCount = reader.ReadUInt16();
        stream.Position += 12;
        var optionalHeaderSize = reader.ReadUInt16();
        stream.Position += 2;

        var optionalHeaderOffset = stream.Position;
        EnsureRange(stream, optionalHeaderOffset, optionalHeaderSize, "optional header");
        var magic = reader.ReadUInt16();
        var isPe32 = magic == Pe32Magic;
        if (!isPe32)
        {
            throw new InvalidDataException($"Expected PE32 optional header 0x10b, found 0x{magic:x4}.");
        }

        const int pe32DataDirectoryOffset = 96;
        if (optionalHeaderSize < pe32DataDirectoryOffset + 8)
        {
            throw new InvalidDataException("PE optional header has no export directory entry.");
        }

        stream.Position = optionalHeaderOffset + pe32DataDirectoryOffset;
        var exportDirectoryRva = reader.ReadUInt32();
        var exportDirectorySize = reader.ReadUInt32();
        if (exportDirectoryRva == 0 || exportDirectorySize == 0)
        {
            throw new InvalidDataException("PE file has no export directory.");
        }

        stream.Position = optionalHeaderOffset + optionalHeaderSize;
        var sections = new List<Section>(sectionCount);
        for (var index = 0; index < sectionCount; index++)
        {
            EnsureRange(stream, stream.Position, 40, "section header");
            stream.Position += 8;
            var virtualSize = reader.ReadUInt32();
            var virtualAddress = reader.ReadUInt32();
            var rawDataSize = reader.ReadUInt32();
            var rawDataPointer = reader.ReadUInt32();
            stream.Position += 16;
            sections.Add(new Section(virtualAddress, virtualSize, rawDataPointer, rawDataSize));
        }

        var exportOffset = RvaToOffset(exportDirectoryRva, sections);
        EnsureRange(stream, exportOffset, 40, "export directory");
        stream.Position = exportOffset + 12;
        var moduleNameRva = reader.ReadUInt32();
        stream.Position = exportOffset + 24;
        var numberOfNames = reader.ReadUInt32();
        stream.Position = exportOffset + 32;
        var addressOfNamesRva = reader.ReadUInt32();

        if (numberOfNames > MaximumExportNames)
        {
            throw new InvalidDataException($"Unreasonable export count: {numberOfNames}.");
        }

        var exportedModuleName = ReadAsciiZ(reader, RvaToOffset(moduleNameRva, sections));
        var namesOffset = RvaToOffset(addressOfNamesRva, sections);
        EnsureRange(stream, namesOffset, checked((long)numberOfNames * 4), "export name table");

        var names = new List<string>((int)numberOfNames);
        for (var index = 0; index < numberOfNames; index++)
        {
            stream.Position = namesOffset + (index * 4L);
            var nameRva = reader.ReadUInt32();
            names.Add(ReadAsciiZ(reader, RvaToOffset(nameRva, sections)));
        }

        names.Sort(StringComparer.Ordinal);
        return new PortableExecutableInspection(machine, isPe32, exportedModuleName, names);
    }

    private static long RvaToOffset(uint rva, IEnumerable<Section> sections)
    {
        foreach (var section in sections)
        {
            var span = Math.Max(section.VirtualSize, section.RawDataSize);
            if (rva >= section.VirtualAddress && rva - section.VirtualAddress < span)
            {
                return checked(section.RawDataPointer + (rva - section.VirtualAddress));
            }
        }

        throw new InvalidDataException($"RVA 0x{rva:x8} is outside all PE sections.");
    }

    private static string ReadAsciiZ(BinaryReader reader, long offset)
    {
        EnsureRange(reader.BaseStream, offset, 1, "ASCII string");
        reader.BaseStream.Position = offset;

        var bytes = new List<byte>();
        while (reader.BaseStream.Position < reader.BaseStream.Length)
        {
            var value = reader.ReadByte();
            if (value == 0)
            {
                return Encoding.ASCII.GetString(bytes.ToArray());
            }

            bytes.Add(value);
            if (bytes.Count > 4096)
            {
                throw new InvalidDataException("PE string exceeded the safety limit.");
            }
        }

        throw new InvalidDataException("PE string was not null-terminated.");
    }

    private static void EnsureRange(Stream stream, long offset, long size, string label)
    {
        if (offset < 0 || size < 0 || offset > stream.Length - size)
        {
            throw new InvalidDataException($"{label} lies outside the file.");
        }
    }

    private sealed record Section(
        uint VirtualAddress,
        uint VirtualSize,
        uint RawDataPointer,
        uint RawDataSize);
}