#pragma once
#include <stdio.h>
#include <windows.h>
#include <winnt.h>

// Forward declarations
typedef struct IORING_HANDLE_REF IORING_HANDLE_REF;
typedef struct IORING_BUFFER_REF IORING_BUFFER_REF;
typedef void *HIORING;

/* --------------------- Types declaration --------------------- */
typedef enum IORING_CREATE_ADVISORY_FLAGS {
  IORING_CREATE_ADVISORY_FLAGS_NONE,
  IORING_CREATE_SKIP_BUILDER_PARAM_CHECKS
} IORING_CREATE_ADVISORY_FLAGS;
// Specifies advisory flags for creating an I/O ring with a call to
// CreateIoRing.

typedef enum IORING_CREATE_REQUIRED_FLAGS {
  IORING_CREATE_REQUIRED_FLAGS_NONE
} IORING_CREATE_REQUIRED_FLAGS;
// Specifies required flags for creating an I/O ring with a call to
// CreateIoRing.

typedef enum IORING_REF_KIND {
  IORING_REF_RAW = 0,
  IORING_REF_REGISTERED = 1,
} IORING_REF_KIND;
// Specifies the type of an IORING_HANDLE_REF structure.

typedef enum IORING_SQE_FLAGS {
  IOSQE_FLAGS_NONE,
  IOSQE_FLAGS_DRAIN_PRECEDING_OPS
} IORING_SQE_FLAGS;
// Specifies kernel behavior options for I/O ring submission queue entries

// IORING_REGISTERED_BUFFER structure
typedef struct IORING_REGISTERED_BUFFER {
  UINT32 Index;
  UINT32 Offset;
} IORING_REGISTERED_BUFFER;

// IORING_HANDLE_REF
struct IORING_HANDLE_REF {
  IORING_REF_KIND Kind;
  union {
    HANDLE Handle;
    UINT32 Index;
  } HandleUnion;
};
// Represents a reference to a file handle used in an I/O ring operation

// IORING_BUFFER_REF
struct IORING_BUFFER_REF {
  IORING_REF_KIND Kind;
  union {
    void *Address;
    IORING_REGISTERED_BUFFER IndexAndOffset;
  } BufferUnion;
};

typedef struct IORING_BUFFER_INFO {
  void *Address;
  UINT32 Length;
} IORING_BUFFER_INFO;

// IORING_BUFFER_REF represents a reference to a buffer used in an I/O ring
// operation

// IORING_VERSION enumeration
typedef enum IORING_VERSION {
  IORING_VERSION_INVALID = 0,
  IORING_VERSION_1 = 1,
  IORING_VERSION_2 = 2,
  IORING_VERSION_3 = 3,
  IORING_VERSION_4 = 4,
} IORING_VERSION;

typedef enum IORING_FEATURE_FLAGS {
  IORING_FEATURE_FLAGS_NONE = 0,
  IORING_FEATURE_UM_EMULATION = 1
} IORING_FEATURE_FLAGS;

// IORING_CAPABILITIES structure
typedef struct IORING_CAPABILITIES {
  IORING_VERSION MaxVersion;
  UINT32 MaxSubmissionQueueSize;
  UINT32 MaxCompletionQueueSize;
  IORING_FEATURE_FLAGS FeatureFlags;
} IORING_CAPABILITIES;
// Represents the IORING API capabilities.

// IORING_CQE structure
typedef struct IORING_CQE {
  UINT_PTR UserData;
  HRESULT ResultCode;
  ULONG_PTR Information;
} IORING_CQE;
// Represents a completed I/O ring queue entry.

// IORING_CREATE_FLAGS structure
typedef struct IORING_CREATE_FLAGS {
  IORING_CREATE_REQUIRED_FLAGS Required;
  IORING_CREATE_ADVISORY_FLAGS Advisory;
} IORING_CREATE_FLAGS;
// Specifies flags for creating an I/O ring with a call to CreateIoRing.

// IORING_INFO structure
typedef struct IORING_INFO {
  IORING_VERSION IoRingVersion;
  IORING_CREATE_FLAGS Flags;
  UINT32 SubmissionQueueSize;
  UINT32 CompletionQueueSize;
} IORING_INFO;
// Represents the shape and version information for the specified I/O ring

// IORING_OP_CODE for IsIoRingOpSupported
typedef enum IORING_OP_CODE {
  IORING_OP_NOP = 0,
  IORING_OP_READ = 1,
  IORING_OP_WRITE = 2,
  IORING_OP_FLUSH = 3,
  IORING_OP_REGISTER_BUFFERS = 4,
  IORING_OP_REGISTER_FILES = 5,
  IORING_OP_CANCEL = 6,
} IORING_OP_CODE;

/* --------------------- Dynamic loader --------------------- */
// Function pointer types
typedef BOOL(WINAPI *IsIoRingOpSupported_t)(HIORING, IORING_OP_CODE);
typedef HRESULT(WINAPI *QueryIoRingCapabilities_t)(IORING_CAPABILITIES *);
typedef HRESULT(WINAPI *GetIoRingInfo_t)(HIORING, IORING_INFO *);
typedef HRESULT(WINAPI *CreateIoRing_t)(IORING_VERSION, IORING_CREATE_FLAGS,
                                        UINT32, UINT32, HIORING *);
typedef HRESULT(WINAPI *CloseIoRing_t)(HIORING);
typedef HRESULT(WINAPI *SubmitIoRing_t)(HIORING, UINT32, UINT32, UINT32 *);
typedef HRESULT(WINAPI *PopIoRingCompletion_t)(HIORING, IORING_CQE *);
typedef HRESULT(WINAPI *SetIoRingCompletionEvent_t)(HIORING, HANDLE);
typedef HRESULT(WINAPI *BuildIoRingCancelRequest_t)(HIORING, IORING_HANDLE_REF,
                                                    UINT_PTR, UINT_PTR);
typedef HRESULT(WINAPI *BuildIoRingReadFile_t)(HIORING, IORING_HANDLE_REF,
                                               IORING_BUFFER_REF, UINT32,
                                               UINT64, UINT_PTR,
                                               IORING_SQE_FLAGS);
typedef HRESULT(WINAPI *BuildIoRingRegisterBuffers_t)(
    HIORING, UINT32, IORING_BUFFER_INFO const[], UINT_PTR);

typedef HRESULT(WINAPI *BuildIoRingRegisterFileHandles_t)(HIORING, UINT32,
                                                          HANDLE const[],
                                                          UINT_PTR);

// Core:
// Queries the support of the specified operation for the specified I/O ring
static IsIoRingOpSupported_t IsIoRingOpSupported = NULL;

// Queries the OS for the supported capabilities for IORINGs
static QueryIoRingCapabilities_t QueryIoRingCapabilities = NULL;

// Gets information about the API version and queue sizes of an I/O ring
static GetIoRingInfo_t GetIoRingInfo = NULL;

// Creates a new instance of an I/O ring submission/completion queue pair and
// returns a handle for referencing the I/O ring
static CreateIoRing_t CreateIoRing = NULL;

// Closes an HIORING handle that was previously opened with a call to
// CreateIoRing
static CloseIoRing_t CloseIoRing = NULL;

// Submission / completion:
// Submits all constructed but not yet submitted entries to the kernel’s queue
// and optionally waits for a set of operations to complete
static SubmitIoRing_t SubmitIoRing = NULL;

// Pops a single entry from the completion queue, if one is available
static PopIoRingCompletion_t PopIoRingCompletion = NULL;

// Registers a completion queue event with an IORING
static SetIoRingCompletionEvent_t SetIoRingCompletionEvent = NULL;

// Operations:
// Performs an asynchronous read from a file using an I/O ring
static BuildIoRingReadFile_t BuildIoRingReadFile = NULL;

// Attempts to cancel a previously submitted I/O ring operation
static BuildIoRingCancelRequest_t BuildIoRingCancelRequest = NULL;

// Registers an array of buffers with the system for future I/O ring operations
static BuildIoRingRegisterBuffers_t BuildIoRingRegisterBuffers = NULL;

// Registers an array of file handles with the system for future I/O ring
// operations
static BuildIoRingRegisterFileHandles_t BuildIoRingRegisterFileHandles = NULL;

static int io_ring_loaded = 0;

static int io_ring_load_functions(void) {
  if (io_ring_loaded)
    return 1;

  HMODULE hKernel = GetModuleHandleW(L"kernel32.dll");
  if (!hKernel)
    return 0;

  IsIoRingOpSupported =
      (IsIoRingOpSupported_t)GetProcAddress(hKernel, "IsIoRingOpSupported");
  QueryIoRingCapabilities = (QueryIoRingCapabilities_t)GetProcAddress(
      hKernel, "QueryIoRingCapabilities");
  GetIoRingInfo = (GetIoRingInfo_t)GetProcAddress(hKernel, "GetIoRingInfo");
  CreateIoRing = (CreateIoRing_t)GetProcAddress(hKernel, "CreateIoRing");
  CloseIoRing = (CloseIoRing_t)GetProcAddress(hKernel, "CloseIoRing");
  SubmitIoRing = (SubmitIoRing_t)GetProcAddress(hKernel, "SubmitIoRing");
  PopIoRingCompletion =
      (PopIoRingCompletion_t)GetProcAddress(hKernel, "PopIoRingCompletion");
  SetIoRingCompletionEvent = (SetIoRingCompletionEvent_t)GetProcAddress(
      hKernel, "SetIoRingCompletionEvent");
  BuildIoRingReadFile =
      (BuildIoRingReadFile_t)GetProcAddress(hKernel, "BuildIoRingReadFile");
  BuildIoRingCancelRequest = (BuildIoRingCancelRequest_t)GetProcAddress(
      hKernel, "BuildIoRingCancelRequest");
  BuildIoRingRegisterBuffers = (BuildIoRingRegisterBuffers_t)GetProcAddress(
      hKernel, "BuildIoRingRegisterBuffers");
  BuildIoRingRegisterFileHandles =
      (BuildIoRingRegisterFileHandles_t)GetProcAddress(
          hKernel, "BuildIoRingRegisterFileHandles");

  io_ring_loaded =
      (IsIoRingOpSupported && QueryIoRingCapabilities && CreateIoRing &&
       CloseIoRing && SubmitIoRing && PopIoRingCompletion &&
       SetIoRingCompletionEvent && BuildIoRingReadFile &&
       BuildIoRingCancelRequest && BuildIoRingRegisterBuffers &&
       BuildIoRingRegisterFileHandles);

  if (io_ring_loaded)
    printf("[I/O Ring] Functions loaded\n");
  else
    printf("[I/O Ring] Some functions not available\n");

  return io_ring_loaded;
}

/* ------------- Standard helper functions definition ------------- */
// Creates an instance of the IORING_BUFFER_REF structure with the provided
// buffer index and offset
static inline IORING_BUFFER_REF
IoRingBufferRefFromIndexAndOffset(UINT32 index, UINT32 offset) {
  IORING_BUFFER_REF ref;
  ref.Kind = IORING_REF_REGISTERED;
  ref.BufferUnion.IndexAndOffset.Index = index;
  ref.BufferUnion.IndexAndOffset.Offset = offset;
  return ref;
}

// Creates an instance of the IORING_BUFFER_REF structure from the provided
// pointer
static IORING_BUFFER_REF IoRingBufferRefFromPointer(void *addr) {
  IORING_BUFFER_REF ref;
  ref.Kind = IORING_REF_RAW;
  ref.BufferUnion.Address = addr;
  return ref;
}

// Creates an instance of the IORING_HANDLE_REF structure from the provided file
// handle
static IORING_HANDLE_REF IoRingHandleRefFromHandle(HANDLE h) {
  IORING_HANDLE_REF ref;
  ref.Kind = IORING_REF_RAW;
  ref.HandleUnion.Handle = h;
  return ref;
}

// Creates an instance of the IORING_HANDLE_REF structure from the provided
// index
static inline IORING_HANDLE_REF IoRingHandleRefFromIndex(UINT32 index) {
  IORING_HANDLE_REF ref;
  ref.Kind = IORING_REF_REGISTERED; // MUST be registered
  ref.HandleUnion.Index = index;
  return ref;
}

// NOTE: If you are using index-based buffers or handles, make sure you have
// successfully called BuildIoRingRegisterBuffers or
// BuildIoRingRegisterFileHandles first so the kernel has a valid table to look
// into, otherwise the kernel will treat the index as an invalid memory
// address/handle.