#ifndef __CAPSULEX64_H__
#define __CAPSULEX64_H__

#include "../uefi_headers/Uefi.h"

#pragma pack(push, 1)

typedef unsigned char       UINT8;
typedef unsigned short      UINT16;
typedef unsigned int        UINT32;
typedef unsigned long long  UINT64;
typedef signed char         INT8;
typedef short               INT16;
typedef int                 INT32;
typedef long long           INT64;
typedef long long           INTN;
typedef unsigned long long  UINTN;
typedef unsigned char       BOOLEAN;
typedef unsigned long long  EFI_STATUS;
typedef void *              EFI_HANDLE;
typedef UINT64              EFI_PHYSICAL_ADDRESS;
typedef UINT64              EFI_VIRTUAL_ADDRESS;
typedef struct __va_list_tag va_list;

#define TRUE                            1
#define FALSE                           0

#define EFI_SUCCESS                    0ULL
#define EFI_INVALID_PARAMETER          0x8000000000000002ULL
#define EFI_BUFFER_TOO_SMALL           0x8000000000000005ULL
#define EFI_NOT_READY                  0x8000000000000006ULL
#define EFI_OUT_OF_RESOURCES           0x8000000000000009ULL
#define EFI_NOT_FOUND                  0x800000000000000EULL

typedef struct {
    UINT32  Data1;
    UINT16  Data2;
    UINT16  Data3;
    UINT8   Data4[8];
} EFI_GUID;

typedef struct {
    UINT64  Signature;
    UINT32  Revision;
    UINT32  HeaderSize;
    UINT32  CRC32;
    UINT32  Reserved;
} EFI_TABLE_HEADER;

typedef struct {
    EFI_TABLE_HEADER    Hdr;
    UINT16              *FirmwareVendor;
    UINT32              FirmwareRevision;
    UINT32              ConsoleInHandle;
    void                *ConIn;
    UINT32              ConsoleOutHandle;
    void                *ConOut;
    UINT32              StandardErrorHandle;
    void                *StdErr;
    void                *RuntimeServices;
    void                *BootServices;
    UINT64              NumberOfTableEntries;
    void                *ConfigurationTable;
} EFI_SYSTEM_TABLE;

typedef struct {
    UINT16  Limit;
    UINT64  Base;
} IDTR;

typedef struct {
    UINT16  Limit;
    UINT64  Base;
} GDTR;

typedef struct {
    EFI_GUID    CapsuleGuid;
    UINT32      HeaderSize;
    UINT32      Flags;
    UINT32      CapsuleImageSize;
} EFI_CAPSULE_HEADER;

typedef struct {
    UINT64  Length;
    union {
        UINT64  DataBlock;
        UINT64  ContinuationPointer;
    } Union;
} EFI_CAPSULE_BLOCK_DESCRIPTOR;

typedef struct {
    UINT64  BaseAddress;
    UINT64  Length;
} MEMORY_RESOURCE_DESCRIPTOR;

typedef struct {
    UINT64  PageTableStack[8];
    UINT64  PageTableStackIndex;
    UINT64  PageMapLevel4Table;
    UINT64  PhysicalAddressBits;
} PAGE_TABLE_CONTEXT;

EFI_STATUS  ModuleEntryPoint(EFI_HANDLE ImageHandle,
                             EFI_SYSTEM_TABLE *SystemTable);
void        CapsuleX64Entry(UINT64 a1, UINT64 a2, UINT64 a3, UINT64 a4);
void        PageFaultHandlerWrapper(UINT64 a1, UINT64 a2,
                                    UINT64 a3, UINT64 a4);
UINT64      PageFaultHandlerCore(void);
void        X64ExceptionHandler(UINT64 ExceptionType, UINT64 *ContextRecord);
void        SetPageFaultHandled(UINT8 Flag);
void        InitPageFaultHandler(void (**Table)());
UINT64      AsmCpuid(UINT32 Leaf, UINT32 *pEax, UINT32 *pEbx,
                     UINT32 *pEcx, UINT32 *pEdx);
UINT64      AsmCpuidEx(UINT32 Leaf, UINT32 Subleaf,
                       UINT32 *pEax, UINT32 *pEbx,
                       UINT32 *pEcx, UINT32 *pEdx);
UINT64      AsmReadCr3(void);
UINT64      AsmReadCr2(void);
UINT16      AsmReadCs(void);
void        AsmReadIdtr(void *pIdtr);
void        AsmWriteIdtr(UINT16 *pLimit);
void        AsmDisablePaging64(UINT16 Signature, UINT64 EntryPoint,
                               EFI_HANDLE ImageHandle,
                               EFI_SYSTEM_TABLE *SystemTable,
                               UINT32 NewStack);
UINT64      ValidateCapsuleIntegrity(UINT64 BlockList,
                        MEMORY_RESOURCE_DESCRIPTOR *MemoryResources);
EFI_STATUS  GetCapsuleInfo(EFI_CAPSULE_BLOCK_DESCRIPTOR *Desc,
                           UINT64 *CapsuleCount, UINT64 *TotalLength,
                           UINT64 *CapsuleNumber);
EFI_STATUS  CapsuleCoalesce(UINT64 ScratchSize,
                            UINT64 *BlockListBuffer,
                            UINT64 *MemoryResource,
                            UINT64 **BlockDescriptors,
                            UINT32 *CapsuleDescriptorCount);
UINT64      AllocatePageTablePage(PAGE_TABLE_CONTEXT *ptCtx,
                                  UINT64 *pPteAddr);
UINT64      BitFieldReadU64(UINT64 Value, UINT64 StartBit, UINT64 EndBit);
UINT64      InternalLShiftU64(UINT64 Value, UINT64 Count);
void        *InternalZeroMem(void *Buffer, UINT64 Size);
void        *InternalCopyMem(void *Dest, const void *Src, UINT64 Count);
void        *ZeroMem(void *Buffer, UINT64 Size);
void        *CopyMem(void *Dest, const void *Src, UINT64 Count);
UINT64      *RelocateCapsuleDescriptors(
                EFI_CAPSULE_BLOCK_DESCRIPTOR *Desc,
                UINT64 *FlatBuffer,
                UINT64 DescriptorCount,
                UINT64 *ScratchBuffer,
                UINT64 ScratchSize);
UINT64      *FindFreeMemoryInDescriptors(UINT64 *ResourceArray,
                                         UINT64 *Target,
                                         UINT64 Size,
                                         UINT64 SourceBase);
BOOLEAN     ValidateMemoryResource(MEMORY_RESOURCE_DESCRIPTOR *MemResources,
                                   UINT64 Address, UINT64 Size);
UINT8       DebugPrint(UINT64 ErrorLevel, const char *Format, ...);
UINT64      AssertFail(UINT64 Filename, INT32 LineNumber,
                       UINT64 Expression);
UINT64      DebugPrintWorker(const char *Format, ...);

/* ---- String helpers ---- */
UINT64      AsciiStrLen(const char *String);
UINT64      InternalVSPrint(UINT8 *Buffer, UINT64 BufferSize,
                            UINT64 Flags, const char *Format,
                            va_list VaList);
UINT8       *InternalPrintPad(UINT8 *Buffer, UINT64 BufferSize,
                              UINT64 Count, UINT16 Char,
                              UINT64 CharWidth);

/* ---- Serial ---- */
UINT64      SerialPortWrite(UINT8 *Buffer, UINT64 Size);

/* ---- PE/COFF ---- */
char        *PeCoffGetPdbPath(char *Pe32Data);

/* ---- Image info ---- */
void        FindAndLogImageInfo(void *ExceptionType, UINT64 *ContextRecord);
UINT64      DumpExceptionContext(void *ExceptionType, UINT64 *ContextRecord);

/* ---- APIC ---- */
UINT64      GetApicId(void);

#pragma pack(pop)

#endif /* __CAPSULEX64_H__ */