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

namespace WhiteMagic.ProcessEnvironment;

/// <summary>
/// Managed reader for a target process's Process Environment Block (PEB).
/// </summary>
public sealed class ManagedPeb
{
    private readonly MemoryBase _memory;
    private readonly IntPtr _pebAddress;

    /// <summary>
    /// Creates a PEB reader for the process associated with the specified memory facade.
    /// </summary>
    public ManagedPeb(MemoryBase memory)
    {
        _memory = memory ?? throw new ArgumentNullException(nameof(memory));
        _pebAddress = QueryPebAddress();
    }

    /// <summary>Returns the native address of the PEB in the target process.</summary>
    public IntPtr ReadPebAddress() => _pebAddress;

    /// <summary>Reads the BeingDebugged byte from the PEB.</summary>
    public byte ReadBeingDebugged()
    {
        return _memory.Read<byte>(_pebAddress + 2);
    }

    /// <summary>Reads the ImageBaseAddress pointer from the PEB.</summary>
    public IntPtr ReadImageBaseAddress()
    {
        int offset = _memory.Is64Bit ? 0x10 : 0x08;
        return ReadPointer(offset);
    }

    /// <summary>Reads the PEB_LDR_DATA pointer from the PEB.</summary>
    public IntPtr ReadLdrAddress()
    {
        int offset = _memory.Is64Bit ? 0x18 : 0x0C;
        return ReadPointer(offset);
    }

    /// <summary>
    /// Determines whether the target process is running under WOW64.
    /// </summary>
    public bool ReadIsWow64Process()
    {
        if (!NativeMethods.IsWow64Process(_memory.Handle, out bool wow64))
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException($"IsWow64Process failed with error {error}.");
        }

        return wow64;
    }

    private IntPtr QueryPebAddress()
    {
        var info = new ProcessBasicInformation();
        int status = NativeMethods.NtQueryInformationProcess(
            _memory.Handle,
            0,
            ref info,
            (uint)Marshal.SizeOf<ProcessBasicInformation>(),
            out _);

        if (status < 0 || info.PebBaseAddress == IntPtr.Zero)
        {
            throw new InvalidOperationException(
                $"NtQueryInformationProcess failed to retrieve the PEB (NTSTATUS {status:X8}).");
        }

        return info.PebBaseAddress;
    }

    private IntPtr ReadPointer(int offset)
    {
        IntPtr address = _pebAddress + offset;

        if (_memory.Is64Bit)
        {
            ulong raw = _memory.Read<ulong>(address);
            return new IntPtr((long)raw);
        }

        uint raw32 = _memory.Read<uint>(address);
        return new IntPtr((int)raw32);
    }
}