/*
 * main.c – UEFI boot loader.
 *
 * This is the first code that runs.  efi_main() is called by UEFI
 * firmware via the PE32+ entry point.  It:
 *   1. Reads kernel.elf from the EFI System Partition
 *   2. Parses the ELF64 headers and maps PT_LOAD segments into memory
 *   3. Populates a BootInfo struct with generic service wrappers
 *   4. Jumps to the kernel entry point (kmain)
 */

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

#include "boot_info.h"

/* ================================================================
 *  ELF64 constants and types
 * ================================================================ */

#define ELF_MAGIC  0x464c457f   /* "\x7fELF" as a little-endian UINT32 */
#define PT_LOAD    1            /* loadable program segment             */

/* Minimal ELF64 file header (enough to locate program headers). */
typedef struct {
    UINT32 e_magic;         /* must be ELF_MAGIC                     */
    UINT8  e_class;         /* 2 = 64-bit                            */
    UINT8  e_data;          /* 1 = little-endian                     */
    UINT8  e_version;
    UINT8  e_osabi;
    UINT8  e_abiversion;
    UINT8  e_pad[7];
    UINT16 e_type;          /* ET_EXEC = 2                           */
    UINT16 e_machine;       /* EM_X86_64 = 62                        */
    UINT32 e_version2;
    UINT64 e_entry;         /* virtual address of entry point        */
    UINT64 e_phoff;         /* file offset to program header table   */
    UINT64 e_shoff;
    UINT32 e_flags;
    UINT16 e_ehsize;
    UINT16 e_phentsize;     /* size of one program header entry      */
    UINT16 e_phnum;         /* number of program header entries      */
    UINT16 e_shentsize;
    UINT16 e_shnum;
    UINT16 e_shstrndx;
} Elf64_Ehdr;

/* ELF64 program header (describes one segment). */
typedef struct {
    UINT32 p_type;          /* segment type (PT_LOAD, etc.)          */
    UINT32 p_flags;         /* segment flags (R/W/X)                 */
    UINT64 p_offset;        /* file offset of segment data           */
    UINT64 p_vaddr;         /* virtual address to map at             */
    UINT64 p_paddr;         /* physical address (usually == vaddr)   */
    UINT64 p_filesz;        /* bytes of segment data in file         */
    UINT64 p_memsz;         /* total bytes in memory (>= filesz)     */
    UINT64 p_align;         /* alignment requirement                 */
} Elf64_Phdr;

/* ================================================================
 *  EFI file-system helpers
 * ================================================================ */

/*
 * Open the root directory of the volume from which this loader was
 * loaded (i.e. the EFI System Partition).
 */
static EFI_STATUS open_root_volume(EFI_HANDLE ImageHandle, EFI_FILE_PROTOCOL **Root)
{
    EFI_STATUS Status;
    EFI_LOADED_IMAGE *LoadedImage = NULL;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *FileSystem = NULL;

    Status = uefi_call_wrapper(BS->HandleProtocol, 3, ImageHandle,
                               &gEfiLoadedImageProtocolGuid, (VOID **)&LoadedImage);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(BS->HandleProtocol, 3, LoadedImage->DeviceHandle,
                               &gEfiSimpleFileSystemProtocolGuid, (VOID **)&FileSystem);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    return uefi_call_wrapper(FileSystem->OpenVolume, 2, FileSystem, Root);
}

/*
 * Read an entire file from the EFI System Partition into a pool buffer.
 * Caller must free *Buffer with BS->FreePool() when done.
 */
static EFI_STATUS read_file_to_buffer(EFI_HANDLE ImageHandle, CHAR16 *Path,
                                      VOID **Buffer, UINTN *Size)
{
    EFI_STATUS Status;
    EFI_FILE_PROTOCOL *Root = NULL;
    EFI_FILE_PROTOCOL *File = NULL;
    EFI_FILE_INFO *Info = NULL;
    UINTN InfoSize = 0;
    UINTN FileSize = 0;

    Status = open_root_volume(ImageHandle, &Root);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(Root->Open, 5, Root, &File, Path,
                               EFI_FILE_MODE_READ, 0);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(File->GetInfo, 4, File, &gEfiFileInfoGuid,
                               &InfoSize, NULL);
    if (Status != EFI_BUFFER_TOO_SMALL) {
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData,
                               InfoSize, (VOID **)&Info);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(File->GetInfo, 4, File, &gEfiFileInfoGuid,
                               &InfoSize, Info);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    FileSize = (UINTN)Info->FileSize;
    Status = uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData,
                               FileSize, Buffer);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(File->Read, 3, File, &FileSize, *Buffer);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    uefi_call_wrapper(BS->FreePool, 1, Info);
    uefi_call_wrapper(File->Close, 1, File);

    *Size = FileSize;
    return EFI_SUCCESS;
}

/* ================================================================
 *  ELF64 kernel loader
 * ================================================================ */

/*
 * Parse an ELF64 image in memory and load all PT_LOAD segments to
 * their specified virtual (== physical) addresses.  On success,
 * *EntryOut is set to the kernel entry point.
 */
static EFI_STATUS load_elf_kernel(VOID *Image, UINTN Size, UINT64 *EntryOut)
{
    Elf64_Ehdr *Ehdr = (Elf64_Ehdr *)Image;
    Elf64_Phdr *Phdr = NULL;
    UINT16 Index = 0;

    if (Size < sizeof(Elf64_Ehdr)) {
        return EFI_LOAD_ERROR;
    }

    if (Ehdr->e_magic != ELF_MAGIC || Ehdr->e_class != 2 || Ehdr->e_data != 1) {
        return EFI_LOAD_ERROR;
    }

    if ((Ehdr->e_phoff + (UINT64)Ehdr->e_phnum * Ehdr->e_phentsize) > Size) {
        return EFI_LOAD_ERROR;
    }

    Phdr = (Elf64_Phdr *)((UINT8 *)Image + Ehdr->e_phoff);

    for (Index = 0; Index < Ehdr->e_phnum; Index++) {
        Elf64_Phdr *Segment = (Elf64_Phdr *)((UINT8 *)Phdr + (Index * Ehdr->e_phentsize));
        UINT64 SegmentStart = 0;
        UINT64 SegmentEnd = 0;
        UINT64 SegmentPages = 0;
        EFI_PHYSICAL_ADDRESS Address = 0;

        if (Segment->p_type != PT_LOAD || Segment->p_memsz == 0) {
            continue;
        }

        if ((Segment->p_offset + Segment->p_filesz) > Size) {
            return EFI_LOAD_ERROR;
        }

        SegmentStart = Segment->p_vaddr & ~0xFFFULL;
        SegmentEnd = (Segment->p_vaddr + Segment->p_memsz + 0xFFFULL) & ~0xFFFULL;
        SegmentPages = (SegmentEnd - SegmentStart) >> 12;

        Address = (EFI_PHYSICAL_ADDRESS)SegmentStart;
        if (EFI_ERROR(uefi_call_wrapper(BS->AllocatePages, 4, AllocateAddress,
                                        EfiLoaderData, SegmentPages, &Address))) {
            return EFI_OUT_OF_RESOURCES;
        }

        CopyMem((VOID *)(UINTN)Segment->p_vaddr,
                (UINT8 *)Image + Segment->p_offset,
                (UINTN)Segment->p_filesz);
        if (Segment->p_memsz > Segment->p_filesz) {
            SetMem((VOID *)(UINTN)(Segment->p_vaddr + Segment->p_filesz),
                   (UINTN)(Segment->p_memsz - Segment->p_filesz), 0);
        }
    }

    *EntryOut = Ehdr->e_entry;
    return EFI_SUCCESS;
}

/* ================================================================
 *  UEFI service wrappers  (passed to the kernel via BootInfo)
 * ================================================================ */

static KSTATUS loader_clear_screen(void)
{
    EFI_STATUS Status = uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
    return (KSTATUS)Status;
}

static KSTATUS loader_set_attribute(UINTN Attribute)
{
    EFI_STATUS Status = uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, Attribute);
    return (KSTATUS)Status;
}

static KSTATUS loader_read_key(KeyEvent *Key)
{
    EFI_STATUS Status;
    EFI_INPUT_KEY EfiKey;
    UINTN Index = 0;

    uefi_call_wrapper(BS->WaitForEvent, 3, 1, &ST->ConIn->WaitForKey, &Index);
    Status = uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, &EfiKey);
    if (!EFI_ERROR(Status) && Key != NULL) {
        Key->scan_code   = EfiKey.ScanCode;
        Key->unicode_char = EfiKey.UnicodeChar;
    }
    return (KSTATUS)Status;
}

static KSTATUS loader_try_read_key(KeyEvent *Key)
{
    EFI_STATUS Status;
    EFI_INPUT_KEY EfiKey;

    Status = uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, &EfiKey);
    if (!EFI_ERROR(Status) && Key != NULL) {
        Key->scan_code   = EfiKey.ScanCode;
        Key->unicode_char = EfiKey.UnicodeChar;
    }
    return (KSTATUS)Status;
}

static void loader_shutdown(void)
{
    uefi_call_wrapper(RT->ResetSystem, 4, EfiResetShutdown, EFI_SUCCESS, 0, NULL);
}

static KSTATUS loader_alloc_pages(UINTN pages, UINT64 *addr)
{
    EFI_PHYSICAL_ADDRESS Phys = 0;
    EFI_STATUS Status = uefi_call_wrapper(BS->AllocatePages, 4,
                                          AllocateAnyPages, EfiLoaderData,
                                          pages, &Phys);
    if (!EFI_ERROR(Status) && addr != NULL) {
        *addr = (UINT64)Phys;
    }
    return (KSTATUS)Status;
}

static KSTATUS loader_free_pages(UINT64 addr, UINTN pages)
{
    EFI_STATUS Status = uefi_call_wrapper(BS->FreePages, 2,
                                          (EFI_PHYSICAL_ADDRESS)addr, pages);
    return (KSTATUS)Status;
}

/* ================================================================
 *  UEFI application entry point
 * ================================================================ */

EFI_STATUS
EFIAPI
efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS Status;
    VOID *KernelImage = NULL;
    UINTN KernelSize = 0;
    UINT64 KernelEntry = 0;
    BootInfo Boot;
    KernelEntryFn EntryFn = NULL;

    /* Initialise the GNU-EFI library */
    InitializeLib(ImageHandle, SystemTable);

    Print(L"Loading kernel...\n\r");

    Status = read_file_to_buffer(ImageHandle, L"\\kernel.elf", &KernelImage, &KernelSize);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to read kernel.elf: %r\n\r", Status);
        return Status;
    }

    Status = load_elf_kernel(KernelImage, KernelSize, &KernelEntry);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to load kernel.elf: %r\n\r", Status);
        return Status;
    }

    /* Populate the BootInfo struct with generic UEFI-backed services */
    Boot.print          = Print;
    Boot.clear_screen   = loader_clear_screen;
    Boot.set_attribute  = loader_set_attribute;
    Boot.read_key       = loader_read_key;
    Boot.try_read_key   = loader_try_read_key;
    Boot.shutdown       = loader_shutdown;
    Boot.alloc_pages    = loader_alloc_pages;
    Boot.free_pages     = loader_free_pages;
    Boot.firmware_vendor = SystemTable->FirmwareVendor;
    Boot.firmware_major  = (SystemTable->FirmwareRevision >> 16) & 0xFFFF;
    Boot.firmware_minor  = SystemTable->FirmwareRevision & 0xFFFF;

    /* Jump to the kernel – this should not return */
    EntryFn = (KernelEntryFn)(UINTN)KernelEntry;
    EntryFn(&Boot);

    Print(L"Kernel returned. Halting.\n\r");
    return EFI_SUCCESS;
}