#ifndef __CAPSULEX64_H__
#define __

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

/*=============================================================================
 *  CapsuleX64.efi - Capsule PEI X64 Module
 *  MD5:   f20dc2f52ef47aa5bb67065361900b0a
 *  SHA256: f375c4b961d442295e12bd69e7648e2f45cf4a0d5c35f1795b611656107dbfcc
 *
 *  Source: MdeModulePkg/Universal/CapsulePei/X64/
 *  PEI module for x64 capsule update support. Handles page fault
 *  management during long-mode capsule coalescing, capsule descriptor
 *  validation/relocation, and the transition back to 32-bit mode via
 *  DisablePaging64.
 *============================================================================*/

#pragma pack(push, 1)

/*=============================================================================
 *  EFI Basic Types
 *============================================================================*/
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

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

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

/*=============================================================================
 *  EFI_SYSTEM_TABLE
 *============================================================================*/
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;

/*=============================================================================
 *  IDTR / GDTR structures
 *============================================================================*/
typedef struct {
    UINT16  Limit;
    UINT64  Base;
} IDTR;

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

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

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

/*=============================================================================
 *  Memory Resource Descriptor
 *============================================================================*/
typedef struct {
    UINT64  BaseAddress;
    UINT64  Length;
} MEMORY_RESOURCE_DESCRIPTOR;

/*=============================================================================
 *  Page table context (ring buffer for page allocation)
 *============================================================================*/
typedef struct {
    UINT64  PageTableStack[8];     /* Page pool ring buffer entries */
    UINT64  PageTableStackIndex;   /* Current index in ring buffer (0-7) */
    UINT64  PageMapLevel4Table;    /* PML4 address/mask */
    UINT64  PhysicalAddressBits;   /* Bits to mask for physical address */
} PAGE_TABLE_CONTEXT;

/*=============================================================================
 *  Function Prototypes
 *============================================================================*/

/* ---- Entry point ---- */
EFI_STATUS  ModuleEntryPoint(EFI_HANDLE ImageHandle,
                             EFI_SYSTEM_TABLE *SystemTable);
void        CapsuleX64Entry(UINT64 a1, UINT64 a2, UINT64 a3, UINT64 a4);

/* ---- Page fault and exception handling ---- */
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)());

/* ---- CPUID ---- */
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);

/* ---- CR/segment ---- */
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);

/* ---- Capsule operations ---- */
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);

/* ---- Page table management ---- */
UINT64      AllocatePageTablePage(PAGE_TABLE_CONTEXT *ptCtx,
                                  UINT64 *pPteAddr);
UINT64      BitFieldReadU64(UINT64 Value, UINT64 StartBit, UINT64 EndBit);
UINT64      InternalLShiftU64(UINT64 Value, UINT64 Count);

/* ---- Memory ---- */
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);

/* ---- Descriptor relocation ---- */
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);

/* ---- Debug ---- */
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__ */