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

namespace WhiteMagic.ThreadEnvironment;

/// <summary>
/// Managed reader for a target thread's Thread Environment Block (TEB).
/// </summary>
public sealed class ManagedTeb : IDisposable
{
    private readonly MemoryBase _memory;
    private readonly SafeMemoryHandle _threadHandle;
    private readonly IntPtr _tebAddress;
    private bool _disposed;

    /// <summary>
    /// Creates a TEB reader for the specified thread in the process associated
    /// with the provided memory facade.
    /// </summary>
    public ManagedTeb(MemoryBase memory, int threadId)
    {
        _memory = memory ?? throw new ArgumentNullException(nameof(memory));

        _threadHandle = NativeMethods.OpenThread(
            ThreadAccess.QueryInformation,
            false,
            threadId);

        if (_threadHandle.IsInvalid)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException(
                $"OpenThread failed for thread {threadId}: error {error}.");
        }

        _tebAddress = QueryTebAddress();
    }

    /// <summary>Returns the native address of the TEB in the target process.</summary>
    public IntPtr ReadTebAddress() => _tebAddress;

    /// <summary>Reads the stack base pointer stored in the TEB.</summary>
    public IntPtr ReadStackBase()
    {
        int offset = _memory.Is64Bit ? 0x08 : 0x04;
        return ReadPointer(offset);
    }

    /// <summary>Reads the stack limit pointer stored in the TEB.</summary>
    public IntPtr ReadStackLimit()
    {
        int offset = _memory.Is64Bit ? 0x10 : 0x08;
        return ReadPointer(offset);
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;
            _threadHandle.Dispose();
        }
    }

    private IntPtr QueryTebAddress()
    {
        var info = new ThreadBasicInformation();
        int status = NativeMethods.NtQueryInformationThread(
            _threadHandle,
            0,
            ref info,
            (uint)Marshal.SizeOf<ThreadBasicInformation>(),
            out _);

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

        return info.TebBaseAddress;
    }

    private IntPtr ReadPointer(int offset)
    {
        IntPtr address = _tebAddress + 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);
    }
}