#pragma once

#include <Windows.h>
#include <cstdint>

#include "../common/nt_types.hpp"

namespace TartarusGate {

// NT syscall typedefs used by the loader.
using fnNtAllocateVirtualMemory = NTSTATUS(NTAPI*)(
    HANDLE ProcessHandle, PVOID* BaseAddress, ULONG_PTR ZeroBits,
    PSIZE_T RegionSize, ULONG AllocationType, ULONG Protect);

using fnNtProtectVirtualMemory = NTSTATUS(NTAPI*)(
    HANDLE ProcessHandle, PVOID* BaseAddress, PSIZE_T RegionSize,
    ULONG NewProtect, PULONG OldProtect);

using fnNtWriteVirtualMemory = NTSTATUS(NTAPI*)(
    HANDLE ProcessHandle, PVOID BaseAddress, PVOID Buffer,
    SIZE_T NumberOfBytesToWrite, PSIZE_T NumberOfBytesWritten);

using fnNtFreeVirtualMemory = NTSTATUS(NTAPI*)(
    HANDLE ProcessHandle, PVOID* BaseAddress, PSIZE_T RegionSize, ULONG FreeType);

using fnNtCreateSection = NTSTATUS(NTAPI*)(
    PHANDLE SectionHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES_NT ObjectAttributes,
    PLARGE_INTEGER MaximumSize, ULONG SectionPageProtection, ULONG AllocationAttributes,
    HANDLE FileHandle);

using fnNtMapViewOfSection = NTSTATUS(NTAPI*)(
    HANDLE SectionHandle, HANDLE ProcessHandle, PVOID* BaseAddress, ULONG_PTR ZeroBits,
    SIZE_T CommitSize, PLARGE_INTEGER SectionOffset, PSIZE_T ViewSize,
    ULONG InheritDisposition, ULONG AllocationType, ULONG Win32Protect);

using fnNtUnmapViewOfSection = NTSTATUS(NTAPI*)(
    HANDLE ProcessHandle, PVOID BaseAddress);

using fnNtOpenProcess = NTSTATUS(NTAPI*)(
    PHANDLE ProcessHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES_NT ObjectAttributes,
    PCLIENT_ID_NT ClientId);

using fnNtCreateThreadEx = NTSTATUS(NTAPI*)(
    PHANDLE ThreadHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES_NT ObjectAttributes,
    HANDLE ProcessHandle, PVOID StartRoutine, PVOID Argument, ULONG CreateFlags,
    ULONG_PTR ZeroBits, SIZE_T StackSize, SIZE_T MaximumStackSize,
    PVOID AttributeList);

using fnNtClose = NTSTATUS(NTAPI*)(HANDLE Handle);

using fnNtQuerySystemInformation = NTSTATUS(NTAPI*)(
    ULONG SystemInformationClass, PVOID SystemInformation,
    ULONG SystemInformationLength, PULONG ReturnLength);

using fnNtCreateFile = NTSTATUS(NTAPI*)(
    PHANDLE FileHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES_NT ObjectAttributes,
    PIO_STATUS_BLOCK_NT IoStatusBlock, PLARGE_INTEGER AllocationSize, ULONG FileAttributes,
    ULONG ShareAccess, ULONG CreateDisposition, ULONG CreateOptions, PVOID EaBuffer,
    ULONG EaLength);

using fnNtReadFile = NTSTATUS(NTAPI*)(
    HANDLE FileHandle, HANDLE Event, PVOID ApcRoutine, PVOID ApcContext,
    PIO_STATUS_BLOCK_NT IoStatusBlock, PVOID Buffer, ULONG Length,
    PLARGE_INTEGER ByteOffset, PULONG Key);

struct SyscallTable {
    fnNtAllocateVirtualMemory  NtAllocateVirtualMemory  = nullptr;
    fnNtProtectVirtualMemory   NtProtectVirtualMemory   = nullptr;
    fnNtWriteVirtualMemory     NtWriteVirtualMemory     = nullptr;
    fnNtFreeVirtualMemory      NtFreeVirtualMemory      = nullptr;
    fnNtCreateSection          NtCreateSection          = nullptr;
    fnNtMapViewOfSection       NtMapViewOfSection       = nullptr;
    fnNtUnmapViewOfSection      NtUnmapViewOfSection     = nullptr;
    fnNtOpenProcess            NtOpenProcess            = nullptr;
    fnNtCreateThreadEx         NtCreateThreadEx         = nullptr;
    fnNtClose                  NtClose                  = nullptr;
    fnNtQuerySystemInformation NtQuerySystemInformation = nullptr;
    PVOID                      gadget                   = nullptr;
    bool                       initialized              = false;
};

// Tartarus Gate + Halo's Gate SSN resolution by export hash (no API name strings).
std::uint32_t ResolveSSN(HMODULE ntdll, std::uint32_t exportHash);

// Find `syscall; ret` gadget in ntdll .text for indirect syscalls.
PVOID FindSyscallGadget(HMODULE ntdll);

// Build indirect syscall stub: mov r10,rcx / mov eax,SSN / jmp [gadget].
PVOID CreateIndirectStub(std::uint32_t ssn, PVOID gadget);

// Initialize full syscall table with indirect stubs.
bool Initialize(SyscallTable& table);

void CleanupStubs();

} // namespace TartarusGate