using System.Buffers.Binary;
using MxGateway.Server.Configuration;

namespace MxGateway.Server.Workers;

internal static class WorkerExecutableValidator
{
    private const ushort ImageFileMachineI386 = 0x014c;
    private const ushort ImageFileMachineAmd64 = 0x8664;
    private const int DosHeaderSignatureOffset = 0;
    private const int PeHeaderOffsetPointer = 0x3c;
    private const int PeSignatureSize = 4;
    private const int MachineOffsetFromPeHeader = PeSignatureSize;
    private const int MinimumHeaderSize = 0x40;

    public static void Validate(
        string executablePath,
        WorkerArchitecture requiredArchitecture)
    {
        ushort machine = ReadMachineType(executablePath);
        ushort expectedMachine = requiredArchitecture switch
        {
            WorkerArchitecture.X86 => ImageFileMachineI386,
            WorkerArchitecture.X64 => ImageFileMachineAmd64,
            _ => throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidExecutable,
                "Worker executable required architecture is unsupported."),
        };

        if (machine != expectedMachine)
        {
            throw new WorkerProcessLaunchException(
                WorkerProcessLaunchErrorCode.InvalidExecutable,
                $"Worker executable architecture does not match required {requiredArchitecture} architecture.");
        }
    }

    private static ushort ReadMachineType(string executablePath)
    {
        byte[] header = new byte[MinimumHeaderSize];
        using FileStream stream = File.OpenRead(executablePath);
        if (stream.Read(header) < header.Length)
        {
            throw InvalidExecutable("Worker executable is too small to contain a valid PE header.");
        }

        if (header[DosHeaderSignatureOffset] != 'M' || header[DosHeaderSignatureOffset + 1] != 'Z')
        {
            throw InvalidExecutable("Worker executable does not contain an MZ header.");
        }

        int peHeaderOffset = BinaryPrimitives.ReadInt32LittleEndian(header.AsSpan(PeHeaderOffsetPointer, sizeof(int)));
        if (peHeaderOffset < MinimumHeaderSize)
        {
            throw InvalidExecutable("Worker executable PE header offset is invalid.");
        }

        byte[] peHeaderBytes = new byte[PeSignatureSize + sizeof(ushort)];
        stream.Position = peHeaderOffset;
        if (stream.Read(peHeaderBytes) < peHeaderBytes.Length)
        {
            throw InvalidExecutable("Worker executable PE header is missing.");
        }

        if (peHeaderBytes[0] != 'P' || peHeaderBytes[1] != 'E' || peHeaderBytes[2] != 0 || peHeaderBytes[3] != 0)
        {
            throw InvalidExecutable("Worker executable does not contain a PE header.");
        }

        return BinaryPrimitives.ReadUInt16LittleEndian(
            peHeaderBytes.AsSpan(MachineOffsetFromPeHeader, sizeof(ushort)));
    }

    private static WorkerProcessLaunchException InvalidExecutable(string message)
    {
        return new WorkerProcessLaunchException(
            WorkerProcessLaunchErrorCode.InvalidExecutable,
            message);
    }
}