using System.Buffers.Binary;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Workers;

namespace ZB.MOM.WW.MxGateway.Tests.Gateway.Workers;

/// <summary>
///     Coverage for <see cref="WorkerExecutableValidator"/> PE-header architecture parsing.
///     The validator reads the DOS <c>MZ</c> stub, follows the PE
///     header offset at <c>0x3c</c>, checks the <c>PE\0\0</c> signature, and compares the
///     machine field against the required <see cref="WorkerArchitecture"/>.
/// </summary>
public sealed class WorkerExecutableValidatorTests : IDisposable
{
    private const ushort ImageFileMachineI386 = 0x014c;
    private const ushort ImageFileMachineAmd64 = 0x8664;

    private readonly List<string> _tempFiles = [];

    /// <summary>Verifies that x86 executable matching required architecture does not throw.</summary>
    [Fact]
    public void Validate_X86ExecutableMatchingRequiredArchitecture_DoesNotThrow()
    {
        string path = WritePeFile(ImageFileMachineI386);

        WorkerExecutableValidator.Validate(path, WorkerArchitecture.X86);
    }

    /// <summary>Verifies that x64 executable matching required architecture does not throw.</summary>
    [Fact]
    public void Validate_X64ExecutableMatchingRequiredArchitecture_DoesNotThrow()
    {
        string path = WritePeFile(ImageFileMachineAmd64);

        WorkerExecutableValidator.Validate(path, WorkerArchitecture.X64);
    }

    /// <summary>Verifies that x64 executable when x86 required throws invalid executable.</summary>
    [Fact]
    public void Validate_X64ExecutableWhenX86Required_ThrowsInvalidExecutable()
    {
        string path = WritePeFile(ImageFileMachineAmd64);

        WorkerProcessLaunchException exception = Assert.Throws<WorkerProcessLaunchException>(
            () => WorkerExecutableValidator.Validate(path, WorkerArchitecture.X86));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
        Assert.Contains("architecture", exception.Message, StringComparison.OrdinalIgnoreCase);
    }

    /// <summary>Verifies that x86 executable when x64 required throws invalid executable.</summary>
    [Fact]
    public void Validate_X86ExecutableWhenX64Required_ThrowsInvalidExecutable()
    {
        string path = WritePeFile(ImageFileMachineI386);

        WorkerProcessLaunchException exception = Assert.Throws<WorkerProcessLaunchException>(
            () => WorkerExecutableValidator.Validate(path, WorkerArchitecture.X64));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
    }

    /// <summary>Verifies that file without MZ header throws invalid executable.</summary>
    [Fact]
    public void Validate_FileWithoutMzHeader_ThrowsInvalidExecutable()
    {
        byte[] bytes = new byte[0x80];
        // Leave the first two bytes as zero so the MZ signature check fails.
        string path = WriteTempFile(bytes);

        WorkerProcessLaunchException exception = Assert.Throws<WorkerProcessLaunchException>(
            () => WorkerExecutableValidator.Validate(path, WorkerArchitecture.X86));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
        Assert.Contains("MZ", exception.Message, StringComparison.Ordinal);
    }

    /// <summary>Verifies that file too small for PE header throws invalid executable.</summary>
    [Fact]
    public void Validate_FileTooSmallForPeHeader_ThrowsInvalidExecutable()
    {
        string path = WriteTempFile([(byte)'M', (byte)'Z']);

        WorkerProcessLaunchException exception = Assert.Throws<WorkerProcessLaunchException>(
            () => WorkerExecutableValidator.Validate(path, WorkerArchitecture.X86));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
    }

    /// <summary>Verifies that file without PE signature throws invalid executable.</summary>
    [Fact]
    public void Validate_FileWithoutPeSignature_ThrowsInvalidExecutable()
    {
        // Build a valid MZ header pointing at a PE offset that holds a wrong signature.
        byte[] bytes = new byte[0x100];
        bytes[0] = (byte)'M';
        bytes[1] = (byte)'Z';
        BinaryPrimitives.WriteInt32LittleEndian(bytes.AsSpan(0x3c, sizeof(int)), 0x80);
        // PE region left as zeros — the "PE\0\0" signature check fails.
        string path = WriteTempFile(bytes);

        WorkerProcessLaunchException exception = Assert.Throws<WorkerProcessLaunchException>(
            () => WorkerExecutableValidator.Validate(path, WorkerArchitecture.X86));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
        Assert.Contains("PE", exception.Message, StringComparison.Ordinal);
    }

    private string WritePeFile(ushort machine)
    {
        const int peHeaderOffset = 0x80;
        byte[] bytes = new byte[peHeaderOffset + 6];
        bytes[0] = (byte)'M';
        bytes[1] = (byte)'Z';
        BinaryPrimitives.WriteInt32LittleEndian(bytes.AsSpan(0x3c, sizeof(int)), peHeaderOffset);
        bytes[peHeaderOffset] = (byte)'P';
        bytes[peHeaderOffset + 1] = (byte)'E';
        bytes[peHeaderOffset + 2] = 0;
        bytes[peHeaderOffset + 3] = 0;
        BinaryPrimitives.WriteUInt16LittleEndian(bytes.AsSpan(peHeaderOffset + 4, sizeof(ushort)), machine);
        return WriteTempFile(bytes);
    }

    private string WriteTempFile(byte[] bytes)
    {
        string path = Path.Combine(Path.GetTempPath(), $"mxgw-pe-{Guid.NewGuid():N}.bin");
        File.WriteAllBytes(path, bytes);
        _tempFiles.Add(path);
        return path;
    }

    /// <summary>Deletes the temporary PE fixture files created by this test class, best-effort.</summary>
    public void Dispose()
    {
        foreach (string path in _tempFiles)
        {
            try
            {
                File.Delete(path);
            }
            catch (IOException)
            {
                // Best-effort cleanup of the temp PE fixtures.
            }
        }

        _tempFiles.Clear();
    }
}
