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

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

struct rsdp {
  CHAR8 sig[8];
  UINT8 Checksum;
  CHAR8 OEMID[6];
  UINT8 Revision;
  UINT32 RsdtAddress;
  UINT32 Len;
  UINT64 XsdtAddress;
  UINT8 ExtChecksum;
  UINT8 reserved[3];
};

struct __attribute((packed)) GraphicConfig {
  UINT64 frame_base;
  UINT64 frame_size;
  UINT64 horizontal_resolution;
  UINT64 vertical_resolution;
  UINT64 pixels_per_scan_line;
};

struct __attribute__((packed)) BootParam {
  UINT64 kernel_entry;
  UINT64 rsdp_addr;
  struct GraphicConfig graphic_config;
//  struct gdt bootstrap_gdt[3];
//  struct gdt_desc bootstrap_gdt_desc;
  UINT64 kernel_addr;
};



struct MemoryMap {
  UINTN BufferSize;
  VOID *Buffer;
  UINTN MapSize;
  UINTN MapKey;
  UINTN DescriptorSize;
  UINT32 DescriptorVersion;
};

void HexDump(UINT8 *bytePtr, UINTN size) {
    for (UINTN i = 0; i < size; i++) {
        if (i % 16 == 0 && i != 0) {
            Print(L"\n");
        }
        Print(L"%02x ", bytePtr[i]);
    }
    Print(L"\n");
}


EFI_STATUS LoadFile(CHAR16 *FileName, EFI_PHYSICAL_ADDRESS *FileAddr, UINTN *FilePageSize) {
    EFI_STATUS Status;
    EFI_FILE_IO_INTERFACE *IOVolume;
    EFI_FILE_HANDLE Root;
    EFI_FILE_HANDLE File;
    Print(L"LoadFile\n");

    // Locate all handles that support the Simple File System Protocol.
    UINTN HandleCount = 5;
    void **HandleBuffer;
    Status = LibLocateHandle(ByProtocol, &FileSystemProtocol, NULL, &HandleCount, &HandleBuffer);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to locate handles for file system protocol: %r\n", Status);
        return Status;
    }

    for (UINTN Index = 0; Index < HandleCount; Index++) {
        // Try to open the file system
        Status = uefi_call_wrapper(BS->HandleProtocol, 3, HandleBuffer[Index], &FileSystemProtocol, (void **)&IOVolume);

        if (!EFI_ERROR(Status)) {
            // Open the root directory
            Status = uefi_call_wrapper(IOVolume->OpenVolume, 2, IOVolume, &Root);
            if (EFI_ERROR(Status)) {
                Print(L"Failed to open root directory: %r\n", Status);
                continue;
            }

            // Now open a file
            Status = uefi_call_wrapper(Root->Open, 5, Root, &File,FileName , EFI_FILE_MODE_READ, 0);
            if (EFI_ERROR(Status)) {
                // Print(L"Failed to open file: %r\n", Status);
                uefi_call_wrapper(Root->Close, 1, Root);
                continue;
            }

            EFI_FILE_INFO *FileInfo;
            UINTN FileInfoSize;
            FileInfoSize = sizeof(EFI_FILE_INFO) + 1024;
            FileInfo = AllocatePool(FileInfoSize);
            Status = uefi_call_wrapper(File->GetInfo, 4, File, &gEfiFileInfoGuid, &FileInfoSize, FileInfo);
            if (EFI_ERROR(Status)) {
                Print(L"Failed to get file info.\n");
                FreePool(FileInfo);
                return Status;
            }

            UINTN FileSize = FileInfo->FileSize;
            Print(L"FileSize = %d\n", FileSize);
            *FilePageSize = (FileSize + 4095) / 4096;
            *FileAddr = 0x40000000;
            Status = uefi_call_wrapper(BS->AllocatePages, 4, AllocateAddress, EfiBootServicesCode, *FilePageSize, FileAddr);
            if (EFI_ERROR(Status)) {
                Print(L"Failed to allocate pages.\n");
                FreePool(FileInfo);
                return Status;
            }

            // skip header
            FileSize = 0x1000;
            Status = uefi_call_wrapper(File->Read, 3, File, &FileSize, (VOID *)*FileAddr);
            if (EFI_ERROR(Status)) {
                Print(L"Failed to load file.\n");
                FreePool(FileInfo);
                return Status;
            }

            // load binary part
            FileSize = FileInfo->FileSize-0x1000;
            Status = uefi_call_wrapper(File->Read, 3, File, &FileSize, (VOID *)*FileAddr);
            if (EFI_ERROR(Status)) {
                Print(L"Failed to load file.\n");
                FreePool(FileInfo);
                return Status;
            }
            Print(L"Successfully loaded file: %s\n", FileName);

            FreePool(FileInfo);
            // Close the file and root directory
            uefi_call_wrapper(File->Close, 1, File);
            uefi_call_wrapper(Root->Close, 1, Root);
        }
    }


    return Status;
}

UINT64 ReadCurrentEL() {
    UINT64 currentEL;
    __asm__ volatile("MRS %0, CurrentEL" : "=r" (currentEL));
    return (currentEL >> 2) & 0x3; // Extract EL value
}

void DumpLoadedImageProtocol(EFI_LOADED_IMAGE *LoadedImage) {
    Print(L"Dumping EFI_LOADED_IMAGE_PROTOCOL\n");
    Print(L"Revision: %x\n", LoadedImage->Revision);
    Print(L"ParentHandle: %x\n", LoadedImage->ParentHandle);
    Print(L"SystemTable: %x\n", LoadedImage->SystemTable);

    Print(L"DeviceHandle: %x\n", LoadedImage->DeviceHandle);
    Print(L"FilePath: %s\n", DevicePathToStr(LoadedImage->FilePath));
    // Note: LoadedImage->Reserved is not dumped as it's meant to be NULL/unused

    Print(L"LoadOptionsSize: %x\n", LoadedImage->LoadOptionsSize);
    Print(L"LoadOptions: %s\n", LoadedImage->LoadOptions);

    Print(L"ImageBase: %x\n", LoadedImage->ImageBase);
    Print(L"ImageSize: %lx\n", LoadedImage->ImageSize);
    Print(L"ImageCodeType: %x\n", LoadedImage->ImageCodeType);
    Print(L"ImageDataType: %x\n", LoadedImage->ImageDataType);

    if (LoadedImage->Unload != NULL) {
        Print(L"Unload: %x\n", LoadedImage->Unload);
    } else {
        Print(L"Unload: NULL\n");
    }
}

EFI_STATUS efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {

    EFI_STATUS Status;
    CHAR16 *KernelFileName = L"\\kernel";
    EFI_PHYSICAL_ADDRESS KernelBaseAddr;

    InitializeLib(ImageHandle, SystemTable);
    Print(L"Initializeing\n");
    UINTN BufferPageSize;

    Print(L"LoadFile: %s\n", KernelFileName);

    Status = LoadFile(KernelFileName, &KernelBaseAddr, &BufferPageSize);
    if (EFI_ERROR(Status)) {
        return Status;
    }
    Print(L"Boot 1\n");

    struct BootParam boot_param;
    boot_param.kernel_addr = KernelBaseAddr ;
    boot_param.kernel_entry = boot_param.kernel_addr ;
    // GetGraphicMode(ImageHandle, &(boot_param.graphic_config));


    Print(L"Boot 2\n");

    void **rsdp = NULL;
    EFI_GUID AcpiTableGuid = ACPI_20_TABLE_GUID; // Use ACPI 2.0 table GUID

    Status = LibGetSystemConfigurationTable(&AcpiTableGuid, (void **)&rsdp);
    if (EFI_ERROR(Status) || rsdp == NULL) {
        Print(L"Error getting ACPI table: %r\n", Status);
        return Status;
    } else {
        boot_param.rsdp_addr = (UINT64)rsdp;
    }

    Print(L"Boot 3\n");

    struct MemoryMap MemoryMap = {4096, NULL, 4096, 0, 0, 0};
    Status = uefi_call_wrapper(SystemTable->BootServices->AllocatePool, 3, EfiLoaderData, MemoryMap.BufferSize, &MemoryMap.Buffer);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to allocate memory to get memory map\n");
        return Status;
    }

    Status = uefi_call_wrapper(SystemTable->BootServices->GetMemoryMap, 5, &MemoryMap.MapSize, (EFI_MEMORY_DESCRIPTOR*)MemoryMap.Buffer, &MemoryMap.MapKey, &MemoryMap.DescriptorSize, &MemoryMap.DescriptorVersion);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to get memory map.\n");
        return Status;
    }
    Print(L"Boot r4\n");

    Print(L"Kernel Entry: %llx\n",boot_param.kernel_entry);
    Print(L"RSDP Address: %llx\n",boot_param.rsdp_addr);
    Print(L"Kernel Address: %llx\n",boot_param.kernel_addr);

    HexDump((UINT8 *)KernelBaseAddr,0x100);

    EFI_DEVICE_PATH *Path;
    EFI_LOADED_IMAGE *LoadedImageParent;
    EFI_LOADED_IMAGE *LoadedImage;
    EFI_HANDLE Image;
    CHAR16 *Options = L"root=/dev/sda2 rootfstype=btrfs rw quiet splash";

    Status = uefi_call_wrapper(BS->OpenProtocol, 6, ImageHandle, &LoadedImageProtocol,(void**)&LoadedImageParent, ImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
    if (EFI_ERROR(Status)) {
        Print(L"Could not get LoadedImageProtocol handler %r\n", Status);
        return Status;
    }
    Print(L"Hello,2!\n");
    if (0) {
    Path = FileDevicePath(LoadedImageParent->DeviceHandle, L"\\kernel");
    if (Path == NULL) {
        Print(L"Could not get device path.");
        return EFI_INVALID_PARAMETER;
    }
    Print(L"Hello,3!\n");
    Status = uefi_call_wrapper(BS->LoadImage, 6, FALSE, ImageHandle, Path, NULL, 0, &Image);
    if (EFI_ERROR(Status)) {
        Print(L"Could not load %r", Status);
        FreePool(Path);
        return Status;
    }
    Print(L"Hello,4!\n");
    Status = uefi_call_wrapper(BS->OpenProtocol, 6, Image, &LoadedImageProtocol, (void**)&LoadedImage, ImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
    if (EFI_ERROR(Status)) {
        Print(L"Could not get LoadedImageProtocol handler %r\n", Status);
        uefi_call_wrapper(BS->UnloadImage, 1, Image);
        FreePool(Path);
        return Status;
    }
    Print(L"Hello,5!\n");
    LoadedImage->LoadOptions = Options;
    LoadedImage->LoadOptionsSize = (StrLen(LoadedImage->LoadOptions)+1) * sizeof(CHAR16);

    DumpLoadedImageProtocol (LoadedImage);

    Status = uefi_call_wrapper(SystemTable->BootServices->ExitBootServices, 2, ImageHandle, MemoryMap.MapKey);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to exit boot services\n");
        return Status;
    }

    // Start the image
    Status = uefi_call_wrapper(SystemTable->BootServices->StartImage, 3, Image,  NULL, NULL);

    uefi_call_wrapper(BS->UnloadImage, 1, Image);
    FreePool(Path);

    } else {

    UINT64 el = ReadCurrentEL();
    Print(L"Current Exception Level: %d\n", el);

    Status = uefi_call_wrapper(SystemTable->BootServices->ExitBootServices, 2, ImageHandle, MemoryMap.MapKey);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to exit boot services\n");
        return Status;

    }


    // Inline assembly to jump to the kernel
    // __asm__ volatile("br %0" : : "r"(0x40000000));

    typedef unsigned long (EntryPoint)(struct BootParam*);
    EntryPoint *Entry = (EntryPoint*)(KernelBaseAddr);
    Entry(&boot_param);

    }

    return Status;
}