// UEFI Bootloader for 64-bit kernel
// GNU-EFI based bootloader - FINAL PRODUCTION VERSION (FULLY COMPATIBLE)
// Complies with UEFI Spec 2.8 + x86_64 System V ABI
// Compatible with old GNU-EFI (<3.0) + C99 compilers + legacy firmware

#include <efi.h>
#include <efilib.h>
#include <stdint.h>
#include <assert.h>

// ==========================
// Compatibility Macros (Critical for Legacy Environments)
// ==========================
// Fix EFI_MEMORY_WB missing in old GNU-EFI versions
#ifndef EFI_MEMORY_WB
    #define EFI_MEMORY_WB 0x0000000000000040ULL  // WriteBack memory attribute
#endif

// Fix _Static_assert missing in C99 (use typedef fallback)
#if __STDC_VERSION__ >= 201112L
    #define STATIC_ASSERT(expr, msg) _Static_assert(expr, #msg)
#else
    #define STATIC_ASSERT(expr, msg) typedef char static_assert_##msg[(expr) ? 1 : -1]
#endif

// Kernel boot info structure
typedef struct {
    unsigned long memmap_addr;
    unsigned long memmap_size;
    unsigned long memmap_desc_size;
    unsigned long kernel_phys_addr;
    unsigned long framebuffer_addr;
    unsigned long framebuffer_size;
    unsigned long framebuffer_width;
    unsigned long framebuffer_height;
    unsigned long framebuffer_pitch;
    unsigned long framebuffer_format;
    unsigned long system_table;
} boot_info_t;

// Kernel entry point type
typedef void (*KernelEntry)(boot_info_t*) __attribute__((noreturn));

// ==========================
// Core Configuration (Type-Safe + Compile-Time Validated)
// ==========================
// Kernel load address (matches kernel link address, 4KB page aligned)
#define KERNEL_LOAD_ADDRESS 0x100000ULL  // unsigned long long (x86_64 safe)
// Fallback address for dynamic allocation (if fixed address is reserved)
#define KERNEL_FALLBACK_ADDRESS 0x200000ULL  // 2MB fallback (avoid 1MB reserved)
// x86_64 code/data alignment requirement (16-byte for ABI compliance)
#define KERNEL_ALIGNMENT 16U

// ==========================
// Stack Configuration (ABI Compliant + Type-Safe)
// ==========================
// x86_64 System V ABI: RSP must be 16-byte aligned BEFORE call instruction
// call pushes 8-byte return address → RSP = STACK_TOP - 8 (still 16-byte aligned)
#define BOOTLOADER_STACK_SIZE 0x20000ULL  // 128KB stack (type-safe)
#define BOOTLOADER_STACK_TOP 0x80000ULL   // 16-byte aligned (0x80000 % 16 = 0)
#define BOOTLOADER_STACK_BOTTOM (BOOTLOADER_STACK_TOP - BOOTLOADER_STACK_SIZE)

// ==========================
// Compile-Time Assertions (Critical Validations)
// ==========================
// Ensure stack top is 16-byte aligned (x86_64 ABI requirement)
STATIC_ASSERT((BOOTLOADER_STACK_TOP % 16U) == 0U, StackTopNotAligned);
// Ensure kernel address is 4KB page aligned (UEFI AllocatePages requirement)
STATIC_ASSERT((KERNEL_LOAD_ADDRESS % 0x1000ULL) == 0ULL, KernelAddrNotPageAligned);
// Ensure stack size is positive and non-zero
STATIC_ASSERT(BOOTLOADER_STACK_SIZE > 0ULL, StackSizeZero);
// Ensure stack region does not overlap with kernel load address
STATIC_ASSERT(BOOTLOADER_STACK_TOP < KERNEL_LOAD_ADDRESS, StackOverlapsKernel);
// Ensure fallback address is page aligned
STATIC_ASSERT((KERNEL_FALLBACK_ADDRESS % 0x1000ULL) == 0ULL, FallbackAddrNotPageAligned);

// ==========================
// Validate Stack Memory Region (Robust Descriptor Traversal)
// ==========================
EFI_STATUS ValidateStackMemory(EFI_SYSTEM_TABLE *SystemTable) {
    EFI_STATUS Status;
    UINTN MemoryMapSize = 0U;
    UINTN MapKey;
    UINTN DescriptorSize;
    UINT32 DescriptorVersion;
    EFI_MEMORY_DESCRIPTOR *MemoryMap = NULL;
    EFI_MEMORY_DESCRIPTOR *Descriptor;
    UINTN NumDescriptors;
    BOOLEAN StackRegionValid = FALSE;

    // Get memory map size (first call - get required size)
    Status = SystemTable->BootServices->GetMemoryMap(
        &MemoryMapSize,
        MemoryMap,
        &MapKey,
        &DescriptorSize,
        &DescriptorVersion
    );
    if (Status != EFI_BUFFER_TOO_SMALL) {
        Print(L"Error: GetMemoryMap failed (size check) - %r\r\n", Status);
        return Status;
    }

    // Allocate buffer with extra space (safety margin for firmware changes)
    MemoryMapSize += 2U * DescriptorSize;
    Status = SystemTable->BootServices->AllocatePool(
        EfiLoaderData,
        MemoryMapSize,
        (VOID**)&MemoryMap
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Allocate MemoryMap failed - %r\r\n", Status);
        return Status;
    }

    // Get actual memory map
    Status = SystemTable->BootServices->GetMemoryMap(
        &MemoryMapSize,
        MemoryMap,
        &MapKey,
        &DescriptorSize,
        &DescriptorVersion
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: GetMemoryMap failed (actual) - %r\r\n", Status);
        SystemTable->BootServices->FreePool(MemoryMap);
        return Status;
    }

    // Check if stack region is usable conventional memory
    // FIX: Traverse ALL descriptors until valid one is found (no premature break)
    Descriptor = MemoryMap;
    NumDescriptors = MemoryMapSize / DescriptorSize;

    for (UINTN i = 0U; i < NumDescriptors; i++) {
        const UINT64 DescStart = Descriptor->PhysicalStart;
        const UINT64 DescEnd = DescStart + (Descriptor->NumberOfPages * 0x1000ULL);

        // Check if descriptor fully covers stack region
        if (DescStart <= BOOTLOADER_STACK_BOTTOM && DescEnd >= BOOTLOADER_STACK_TOP) {
            // Must be writable conventional memory (EfiConventionalMemory + WriteBack)
            if (Descriptor->Type == EfiConventionalMemory &&
                (Descriptor->Attribute & EFI_MEMORY_WB)) {
                StackRegionValid = TRUE;
                break;  // Only break when valid descriptor is found
            }
            // Continue traversal if descriptor covers stack but is invalid
        }

        // Move to next descriptor (correct pointer arithmetic for x86_64)
        Descriptor = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)Descriptor + DescriptorSize);
    }

    // Cleanup memory map (safe to free - only used for validation)
    SystemTable->BootServices->FreePool(MemoryMap);

    // Final validation check
    if (!StackRegionValid) {
        Print(L"Error: Stack region 0x%lx-0x%lx is not usable conventional memory\r\n",
               BOOTLOADER_STACK_BOTTOM, BOOTLOADER_STACK_TOP);
        return EFI_INVALID_PARAMETER;
    }

    Print(L"✅ Stack region validated: 0x%lx-0x%lx\r\n",
           BOOTLOADER_STACK_BOTTOM, BOOTLOADER_STACK_TOP);
    return EFI_SUCCESS;
}

// ==========================
// Read Kernel File from Disk (Resource-Safe + Fallback Allocation)
// ==========================
EFI_STATUS ReadKernelFile(
    EFI_HANDLE ImageHandle,
    EFI_SYSTEM_TABLE *SystemTable,
    VOID **KernelBuffer,
    UINTN *KernelSize
) {
    EFI_STATUS Status;
    EFI_LOADED_IMAGE *LoadedImage = NULL;
    EFI_FILE_IO_INTERFACE *FileIo = NULL;
    EFI_FILE_HANDLE RootDir = NULL;
    EFI_FILE_HANDLE KernelFile = NULL;
    EFI_FILE_INFO *FileInfo = NULL;
    UINTN InfoSize = 0U;
    const CHAR16 KernelPath[] = L"\\sys\\kernel\\x86_64\\kernel.bin";

    // Initialize outputs to safe values (prevent wild pointers)
    *KernelBuffer = NULL;
    *KernelSize = 0U;

    // Get Loaded Image Protocol (to find boot device)
    Status = SystemTable->BootServices->HandleProtocol(
        ImageHandle,
        &gEfiLoadedImageProtocolGuid,
        (VOID**)&LoadedImage
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Get LoadedImage Protocol - %r\r\n", Status);
        return Status;
    }

    // Get Simple File System Protocol
    Status = SystemTable->BootServices->HandleProtocol(
        LoadedImage->DeviceHandle,
        &gEfiSimpleFileSystemProtocolGuid,
        (VOID**)&FileIo
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Get FileSystem Protocol - %r\r\n", Status);
        return Status;
    }

    // Open root directory
    Status = FileIo->OpenVolume(FileIo, &RootDir);
    if (EFI_ERROR(Status)) {
        Print(L"Error: Open Root Directory - %r\r\n", Status);
        return Status;
    }

    // Open kernel file (read-only)
    Status = RootDir->Open(
        RootDir,
        &KernelFile,
        (CHAR16*)KernelPath,
        EFI_FILE_MODE_READ,
        0U
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Open Kernel File %s - %r\r\n", KernelPath, Status);
        RootDir->Close(RootDir);
        return Status;
    }

    // Get file size (first call - get required buffer size)
    Status = KernelFile->GetInfo(
        KernelFile,
        &gEfiFileInfoGuid,
        &InfoSize,
        NULL
    );
    if (Status != EFI_BUFFER_TOO_SMALL) {
        Print(L"Error: Get FileInfo Size - %r\r\n", Status);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return Status;
    }

    // Allocate buffer for file info
    Status = SystemTable->BootServices->AllocatePool(
        EfiLoaderData,
        InfoSize,
        (VOID**)&FileInfo
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Allocate FileInfo - %r\r\n", Status);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return Status;
    }

    // Get actual file info
    Status = KernelFile->GetInfo(
        KernelFile,
        &gEfiFileInfoGuid,
        &InfoSize,
        FileInfo
    );
    if (EFI_ERROR(Status)) {
        Print(L"Error: Get FileInfo - %r\r\n", Status);
        SystemTable->BootServices->FreePool(FileInfo);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return Status;
    }

    // Validate file size (non-empty)
    *KernelSize = FileInfo->FileSize;
    if (*KernelSize == 0U) {
        Print(L"Error: Kernel File is Empty\r\n");
        SystemTable->BootServices->FreePool(FileInfo);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return EFI_INVALID_PARAMETER;
    }

    // Allocate physical pages (fixed address first, fallback to dynamic)
    const UINTN PagesNeeded = (*KernelSize + 0xFFFU) / 0x1000U;  // Round up to 4KB pages
    EFI_PHYSICAL_ADDRESS KernelPhysAddr = KERNEL_LOAD_ADDRESS;

    // Step 1: Try fixed address allocation (preferred)
    Status = SystemTable->BootServices->AllocatePages(
        AllocateAddress,    // Fixed address allocation (matches kernel link address)
        EfiLoaderCode,      // Memory type: executable code
        PagesNeeded,
        &KernelPhysAddr
    );

    // Step 2: Fallback to dynamic allocation if fixed address fails
    if (EFI_ERROR(Status)) {
        Print(L"Warning: Fixed address allocation (0x%lx) failed - %r\r\n",
               KERNEL_LOAD_ADDRESS, Status);
        Print(L"Falling back to dynamic allocation at 0x%lx...\r\n", KERNEL_FALLBACK_ADDRESS);

        KernelPhysAddr = KERNEL_FALLBACK_ADDRESS;
        Status = SystemTable->BootServices->AllocatePages(
            AllocateAnyPages,  // Dynamic allocation (UEFI chooses address)
            EfiLoaderCode,
            PagesNeeded,
            &KernelPhysAddr
        );

        if (EFI_ERROR(Status)) {
            Print(L"Error: Dynamic allocation failed - %r\r\n", Status);
            SystemTable->BootServices->FreePool(FileInfo);
            KernelFile->Close(KernelFile);
            RootDir->Close(RootDir);
            return Status;
        }
    }

    // Validate kernel address alignment (x86_64 requirement)
    if ((UINTN)KernelPhysAddr % KERNEL_ALIGNMENT != 0U) {
        Print(L"Error: Kernel Address Not Aligned (0x%lx, req: %d bytes)\r\n",
               KernelPhysAddr, KERNEL_ALIGNMENT);
        SystemTable->BootServices->FreePages(KernelPhysAddr, PagesNeeded);
        SystemTable->BootServices->FreePool(FileInfo);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return EFI_INVALID_PARAMETER;
    }

    // Read kernel file into allocated memory
    *KernelBuffer = (VOID*)KernelPhysAddr;
    Status = KernelFile->Read(KernelFile, KernelSize, *KernelBuffer);
    if (EFI_ERROR(Status)) {
        Print(L"Error: Read Kernel File - %r\r\n", Status);
        SystemTable->BootServices->FreePages(KernelPhysAddr, PagesNeeded);
        SystemTable->BootServices->FreePool(FileInfo);
        KernelFile->Close(KernelFile);
        RootDir->Close(RootDir);
        return Status;
    }

    // Success - print kernel info
    Print(L"✅ Kernel Loaded Successfully\r\n");
    Print(L"  Address: 0x%lx (aligned to %d bytes)\r\n", KernelPhysAddr, KERNEL_ALIGNMENT);
    Print(L"  Size:    %lu bytes (%lu pages)\r\n", *KernelSize, PagesNeeded);

    // Cleanup resources (all allocated resources freed)
    SystemTable->BootServices->FreePool(FileInfo);
    KernelFile->Close(KernelFile);
    RootDir->Close(RootDir);

    return EFI_SUCCESS;
}

// ==========================
// 修复版：兼容老旧/定制UEFI固件的ExitBootServices + 内核跳转
// 核心：显示所有GOP模式 + 默认使用第0个（第一个）模式
// ==========================
EFI_STATUS ExitBootServicesAndJumpToKernel(
    EFI_HANDLE ImageHandle,
    EFI_SYSTEM_TABLE *SystemTable,
    VOID *KernelBuffer
) {
    EFI_STATUS Status = EFI_SUCCESS;

    // ==========================
    // 前置校验
    // ==========================
    UINT64 ActualKernelAddr = (UINT64)KernelBuffer;
    Print(L"Kernel Address: 0x%lx\r\n", ActualKernelAddr);

    UINT8 *KernelCode = (UINT8*)KernelBuffer;
    Print(L"Kernel First Byte: 0x%02x\r\n", KernelCode[0]);

    // ==========================
    // 获取GOP
    // ==========================
    EFI_GRAPHICS_OUTPUT_PROTOCOL *Gop = NULL;
    EFI_STATUS GopStatus = SystemTable->BootServices->LocateProtocol(
        &gEfiGraphicsOutputProtocolGuid,
        NULL,
        (VOID**)&Gop
    );

    unsigned long framebuffer_addr = 0;
    unsigned long framebuffer_size = 0;
    unsigned long framebuffer_width = 0;
    unsigned long framebuffer_height = 0;
    unsigned long framebuffer_pitch = 0;
    unsigned long framebuffer_format = 0;

    Print(L"GopStatus: %r\r\n", GopStatus);
    Print(L"Gop: 0x%lx\r\n", (unsigned long)Gop);

    // ==========================
    // 显示所有 GOP 模式 + 分辨率，并使用第 0 个模式
    // ==========================
    if (!EFI_ERROR(GopStatus) && Gop != NULL) {
        Print(L"\n=== All Supported Graphics Modes ===\r\n");
        UINTN MaxMode = Gop->Mode->MaxMode;

        // 遍历并打印所有模式
        for (UINTN i = 0; i < MaxMode; i++) {
            EFI_GRAPHICS_OUTPUT_MODE_INFORMATION *Info;
            UINTN SizeOfInfo;
            EFI_STATUS s = Gop->QueryMode(Gop, i, &SizeOfInfo, &Info);
            if (EFI_ERROR(s)) continue;

            UINT32 W = Info->HorizontalResolution;
            UINT32 H = Info->VerticalResolution;
            Print(L"Mode %2lu: %lux%lu\r\n", i, W, H);
        }
        Print(L"====================================\r\n");

        // ==========================
        // 关键：强制使用模式
        // ==========================
        UINTN UseMode = 15;
        Status = Gop->SetMode(Gop, UseMode);
        if (!EFI_ERROR(Status)) {
            Print(L"✅ Using Mode %lu: %lux%lu\r\n",
                UseMode,
                Gop->Mode->Info->HorizontalResolution,
                Gop->Mode->Info->VerticalResolution);
        }

        // 读取帧缓冲信息
        if (Gop->Mode != NULL && Gop->Mode->FrameBufferBase != 0) {
            framebuffer_addr = (unsigned long)Gop->Mode->FrameBufferBase;
            framebuffer_size = Gop->Mode->FrameBufferSize;
            framebuffer_width = Gop->Mode->Info->HorizontalResolution;
            framebuffer_height = Gop->Mode->Info->VerticalResolution;
            framebuffer_pitch = Gop->Mode->Info->PixelsPerScanLine * 4;

            // 像素格式
            switch (Gop->Mode->Info->PixelFormat) {
                case PixelRedGreenBlueReserved8BitPerColor: framebuffer_format = 0; break;
                case PixelBlueGreenRedReserved8BitPerColor: framebuffer_format = 1; break;
                case PixelBitMask: framebuffer_format = 2; break;
                case PixelBltOnly: framebuffer_format = 3; break;
                default: framebuffer_format = 1; break;
            }

            UINT64 calc_size = framebuffer_width * framebuffer_height * 4;
            if (framebuffer_size < calc_size) framebuffer_size = calc_size;
        }
    }

    // ==========================
    // 兜底：GOP不可用时使用 1920×1080
    // ==========================
    if (framebuffer_addr == 0) {
        Print(L"Warning: Using fallback 1920x1080\r\n");
        framebuffer_addr = 0x80000000;
        framebuffer_width  = 1920;
        framebuffer_height = 1080;
        framebuffer_pitch  = 1920 * 4;
        framebuffer_size   = 1920 * 1080 * 4;
        framebuffer_format = 1;
    }

    Print(L"Final FB: 0x%lx | %lux%lu | pitch %lu | format %lu\r\n",
        framebuffer_addr, framebuffer_width, framebuffer_height, framebuffer_pitch, framebuffer_format);

    // ==========================
    // ExitBootServices 流程
    // ==========================
    Print(L"Exiting Boot Services...\r\n");

    UINTN MemoryMapSize = 0;
    UINTN MapKey;
    UINTN DescriptorSize;
    UINT32 DescriptorVersion;
    EFI_MEMORY_DESCRIPTOR *MemoryMap = NULL;

    do {
        Status = SystemTable->BootServices->GetMemoryMap(
            &MemoryMapSize, NULL, &MapKey, &DescriptorSize, &DescriptorVersion
        );
        if (Status != EFI_BUFFER_TOO_SMALL) break;

        if (MemoryMap) SystemTable->BootServices->FreePool(MemoryMap);
        MemoryMapSize += 2 * DescriptorSize;

        Status = SystemTable->BootServices->AllocatePool(
            EfiLoaderData, MemoryMapSize, (VOID**)&MemoryMap
        );
        if (EFI_ERROR(Status)) return Status;

        Status = SystemTable->BootServices->GetMemoryMap(
            &MemoryMapSize, MemoryMap, &MapKey, &DescriptorSize, &DescriptorVersion
        );
    } while (Status == EFI_BUFFER_TOO_SMALL);

    if (EFI_ERROR(Status)) return Status;

    // 将MemoryMap复制到静态缓冲区，避免ExitBootServices后数据被覆盖
    static UINT8 SafeMemmap[16384];  // 16KB should be enough for memory map
    if (MemoryMapSize <= sizeof(SafeMemmap)) {
        __builtin_memcpy(SafeMemmap, MemoryMap, MemoryMapSize);
    }

    // 构造启动信息
    boot_info_t BootInfo;
    BootInfo.memmap_addr = (unsigned long)(MemoryMapSize <= sizeof(SafeMemmap) ? (void*)SafeMemmap : (void*)MemoryMap);
    BootInfo.memmap_size = MemoryMapSize;
    BootInfo.memmap_desc_size = DescriptorSize;
    BootInfo.kernel_phys_addr = (unsigned long)KernelBuffer;
    BootInfo.framebuffer_addr = framebuffer_addr;
    BootInfo.framebuffer_size = framebuffer_size;
    BootInfo.framebuffer_width = framebuffer_width;
    BootInfo.framebuffer_height = framebuffer_height;
    BootInfo.framebuffer_pitch = framebuffer_pitch;
    BootInfo.framebuffer_format = framebuffer_format;
    BootInfo.system_table = 0;

    static boot_info_t static_boot_info;
    static_boot_info = BootInfo;
    void* stack_top = (void*)0x80000;

    // 退出UEFI引导服务
    Status = SystemTable->BootServices->ExitBootServices(ImageHandle, MapKey);
    if (EFI_ERROR(Status)) return Status;

    // 跳转到内核
    __asm__ volatile (
        ".intel_syntax noprefix\n"
        "cli\n"
        "cld\n"
        "mov rdi, %1\n"
        "mov rsp, %2\n"
        "jmp %0\n"
        ".att_syntax\n"
        : : "r"(KernelBuffer), "r"(&static_boot_info), "r"(stack_top)
    );

    while (1) __asm__ volatile ("hlt");
    return EFI_SUCCESS;
}

// ==========================
// Main UEFI Entry Point (Clean + Readable)
// ==========================
EFI_STATUS efi_main (EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {
    EFI_STATUS Status;
    VOID *KernelBuffer = NULL;
    UINTN KernelSize = 0U;
    UINT64 CurrentRsp;

    // Initialize GNU-EFI library (required for Print/Protocol access)
    InitializeLib(ImageHandle, SystemTable);

    // Print bootloader header (user feedback)
    Print(L"\n=== UD UEFI Bootloader (x86_64) ===\r\n");
    Print(L"=====================================\r\n");

    // Print initial system info (debug/validation)
    __asm__ volatile ("mov %%rsp, %0" : "=r" (CurrentRsp));
    Print(L"Initial RSP: 0x%lx\r\n", CurrentRsp);
    Print(L"Kernel Load Address (fixed): 0x%lx\r\n", KERNEL_LOAD_ADDRESS);
    Print(L"Kernel Fallback Address: 0x%lx\r\n", KERNEL_FALLBACK_ADDRESS);
    Print(L"Stack Region: 0x%lx-0x%lx\r\n",
           BOOTLOADER_STACK_BOTTOM, BOOTLOADER_STACK_TOP);

    // Step 1: Load kernel from disk
    Print(L"\n[1/3] Loading Kernel...\r\n");
    Status = ReadKernelFile(ImageHandle, SystemTable, &KernelBuffer, &KernelSize);
    if (EFI_ERROR(Status)) {
        Print(L"Error: Kernel Load Failed - %r\r\n", Status);
        return Status;
    }

    // Step 2: Exit boot services and jump to kernel
    Print(L"\n[2/3] Exiting Boot Services...(If you see 3/3, that's mean the kernel has finished, that is an error!)\r\n");
    Status = ExitBootServicesAndJumpToKernel(ImageHandle, SystemTable, KernelBuffer);
    Print(L"\n[3/3] OK...\r\n");
    if (EFI_ERROR(Status)) {
        Print(L"Error: Jump to Kernel Failed - %r\r\n", Status);
        return Status;
    }

    // Unreachable (kernel is marked noreturn)
    return EFI_SUCCESS;
}