/* This software is licensed by the MIT License, see LICENSE file */
/*                           Copyright © 2024-2025 Gregory Lirent */

using System.Runtime.InteropServices;
using System.Runtime.Versioning;
using System.Text.Json.Serialization;
using WebmrAPI.Exceptions;
using WebmrAPI.Models;

namespace WebmrAPI.Utils
{
    [SupportedOSPlatform("windows")]
    public class WindowsProcess : IDisposable
    {
        private IntPtr _ptr = IntPtr.Zero;
        private int _pid = 0;
        private bool _readAccess = false;
        private static int LastError { get => Marshal.GetLastWin32Error(); }
        private static uint MBISize { get => (uint)Marshal.SizeOf(typeof(MemoryBasicInformation)); }

        [Flags]
        [JsonConverter(typeof(JsonEnumConverter<MemoryState>))]
        public enum MemoryState : uint
        {
            Commit = 0x00001000,
            Reserve = 0x00002000,
            Free = 0x00010000,
        }

        [Flags]
        [JsonConverter(typeof(JsonEnumConverter<MemoryType>))]
        public enum MemoryType : uint
        {
            Image = 0x01000000,
            Mapped = 0x00040000,
            Private = 0x00020000,
        }

        [Flags]
        [JsonConverter(typeof(JsonEnumConverter<MemoryPageProtectionState>))]
        public enum MemoryPageProtectionState : uint
        {
            NoAccess = 0x0001,
            ReadOnly = 0x0002,
            ReadWrite = 0x0004,
            WriteCopy = 0x0008,
            Execute = 0x0010,
            ExecuteRead = 0x0020,
            ExecuteReadWrite = 0x0040,
            ExecuteWriteCopy = 0x0080,
            Guard = 0x0100,
            NoCache = 0x0200,
            WriteCombine = 0x0400,
        }

        [Flags]
        public enum ModuleFilterFlags : uint
        {
            Default = 0x00,
            Modules32Bit = 0x01,
            Modules64Bit = 0x02,
            ModulesAll = Modules32Bit | Modules64Bit,
        }

        // https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-memory_basic_information
        [StructLayout(LayoutKind.Sequential)]
        public struct MemoryBasicInformation
        {
            public IntPtr BaseAddress;
            public IntPtr AllocationBase;
            public uint AllocationProtect;
            public ushort PartitionId;
            public UIntPtr RegionSize;
            public MemoryState State;
            public MemoryPageProtectionState PageProtection;
            public MemoryType Type;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct ModuleInformation
        {
            public IntPtr BaseAddress;
            public uint MemorySize;
            public IntPtr EntryPointAddress;
            public System.Text.StringBuilder Name;
            public System.Text.StringBuilder FileName;
        }

        // https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-openprocess
        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern IntPtr OpenProcess(
            uint dwDesiredAccess,
            [MarshalAs(UnmanagedType.Bool)] bool bInheritHandle,
            int dwProcessId
        );

        // https://learn.microsoft.com/en-us/windows/win32/api/handleapi/nf-handleapi-closehandle
        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool CloseHandle(IntPtr hObject);

        // https://learn.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-virtualqueryex
        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern int VirtualQueryEx(
            IntPtr hProcess,
            IntPtr lpAddress,
            out MemoryBasicInformation lpBuffer,
            uint dwLength
        );

        // https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-enumprocessmodulesex
        [DllImport("psapi.dll", SetLastError = true)]
        private static extern bool EnumProcessModulesEx(
            IntPtr hProcess,
            [Out] IntPtr[] lphModule,
            uint cb,
            out uint lpcbNeeded,
            ModuleFilterFlags dwFilterFlag
        );

        // https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-getmodulefilenameexw
        [DllImport("psapi.dll", SetLastError = true, CharSet = CharSet.Unicode)]
        private static extern uint GetModuleFileNameEx(
            IntPtr hProcess,
            IntPtr hModule,
            [Out] System.Text.StringBuilder lpFilename,
            uint nSize
        );

        // https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-getmodulebasenamew
        [DllImport("psapi.dll", SetLastError = true, CharSet = CharSet.Unicode)]
        private static extern uint GetModuleBaseName(
            IntPtr hProcess,
            IntPtr hModule,
            [Out] System.Text.StringBuilder lpBaseName,
            uint nSize
        );

        // https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-getmoduleinformation
        [DllImport("psapi.dll", SetLastError = true)]
        private static extern bool GetModuleInformation(
            IntPtr hProcess,
            IntPtr hModule,
            out ModuleInformation lpmodinfo,
            uint cb
        );

        public void ForeachMemoryRegion(Action<MemoryBasicInformation> callback, MemoryState filter = MemoryState.Commit | MemoryState.Reserve)
        {
            IntPtr cur = IntPtr.Zero;
            long next = 0;
            var regions = new Dictionary<long, MemoryRegionInfo>();

            while ((next > cur.ToInt64() && next < 0x7fffffffffff) || next == 0)
            {
                MemoryBasicInformation info;

                cur = new IntPtr(next);

                if (VirtualQueryEx(_ptr, cur, out info, MBISize) == 0)
                {
                    throw new MemoryRegionException($"WinAPI call 'VirtualQueryEx for PID {_pid} at address 0x{cur.ToInt64():X16}' failed.");
                }

                if ((filter & info.State) != 0)
                {
                    callback(info);
                }
                
                next = info.BaseAddress.ToInt64() + (long)info.RegionSize.ToUInt64();
            }
        }

        public void ForeachLoadedModule(Action<ModuleInformation> callback, ModuleFilterFlags filter = ModuleFilterFlags.ModulesAll, bool force = false, bool ignoreErrors = false)
        {
            if (!_readAccess && force)
            {
                Dispose();
                _readAccess = true;
                Open();
            }

            IntPtr[] modules = new IntPtr[1024];
            uint nmemb;

            if (EnumProcessModulesEx(_ptr, modules, (uint)(modules.Length * IntPtr.Size), out nmemb, filter))
            {
                uint count = nmemb / (uint)IntPtr.Size;
                for (int i = 0; i < count; i++)
                {
                    ModuleInformation info;

                    info.Name = new System.Text.StringBuilder(256);
                    info.FileName = new System.Text.StringBuilder(1024);

                    if ((!GetModuleInformation(_ptr, modules[i], out info, (uint)(2 * IntPtr.Size + sizeof(uint))) ||
                        GetModuleBaseName(_ptr, modules[i], info.Name, (uint)info.Name.Capacity) <= 0 ||
                        GetModuleFileNameEx(_ptr, modules[i], info.FileName, (uint)info.FileName.Capacity) <= 0) && !ignoreErrors)
                    {
                        throw new GettingModuleInfoException($"Failed to get module info for module handle {modules[i]} in PID {_pid}. Error: {LastError}");
                    }

                    callback(info);
                }
            }
            else
            {
                throw new ProcessMonitorException($"EnumProcessModulesEx for PID {_pid} failed. Error: {LastError}");
            }
        }

        private void Open()
        {
            var access = (uint)(_readAccess ? 0x00000010 : 0) | 0x00000400;
            _ptr = OpenProcess(access, false, _pid);

            if (_ptr == IntPtr.Zero)
            {
                var errno = LastError;

                if (errno == 5)
                {
                    throw new ProcessAccessDeniedException($"Access denied to process {_pid}");
                }
                else
                {
                    throw new ProcessMonitorException($"WinAPI call 'OpenProcess for PID {_pid}' failed. Error: {errno}");
                }
            }
        }

        public WindowsProcess(int pid, bool readAccess = false)
        {
            _readAccess = readAccess;
            _pid = pid;
            Open();
        }

        public void Dispose()
        {
            if (_ptr != IntPtr.Zero)
            {
                try { CloseHandle(_ptr); } catch (Exception) { }
            }
        }

        ~WindowsProcess()
        {
            Dispose();
        }
    }
}