/*
  StarForth — Steady-State Virtual Machine Runtime

  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.

  You may obtain a copy of the License at:
      https://github.com/star.4th@proton.me/StarForth/LICENSE.txt

  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose, and noninfringement.

 See the License for the specific language governing permissions and
 limitations under the License.

 StarForth — Steady-State Virtual Machine Runtime
  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.

  You may obtain a copy of the License at:
       https://github.com/star.4th@proton.me/StarForth/LICENSE.txt

  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose, and noninfringement.

  See the License for the specific language governing permissions and
  limitations under the License.

 */

/**
 * uefi_loader.c - UEFI boot loader for StarKernel
 *
 * This loader loads the StarKernel ELF binary from the ESP,
 * parses it, loads segments, applies relocations, and jumps to entry.
 */

#include "uefi.h"
#include "arch.h"
#include "elf_loader.h"
#include "elf64.h"
#include "starkernel/cmdline.h"
#include "starkernel/boot_info_offsets.h"
#include <stddef.h>

/* Verify BootInfo field offsets match boot_info_offsets.h (BOOT_INFO_OFFSETS) */
_Static_assert(offsetof(BootInfo, kernel_stack_base) == BOOT_INFO_KERNEL_STACK_BASE_OFFSET,
    "BOOT_INFO_KERNEL_STACK_BASE_OFFSET mismatch — update boot_info_offsets.h");
_Static_assert(offsetof(BootInfo, kernel_stack_size) == BOOT_INFO_KERNEL_STACK_SIZE_OFFSET,
    "BOOT_INFO_KERNEL_STACK_SIZE_OFFSET mismatch — update boot_info_offsets.h");

/* For monolithic build: kernel_main is linked directly */
#ifdef MONOLITHIC_BUILD
extern void kernel_main(BootInfo *boot_info);
#endif

#if defined(ARCH_AMD64)
#define COM1_BASE 0x3F8
/**
 * @brief Write a byte to an x86 I/O port via @c OUT (early-boot raw path).
 *
 * Used exclusively by the amd64 early-boot serial helper functions
 * (@c raw_serial_init(), @c raw_serial_putc()) before the kernel console
 * subsystem is initialised. Operates identically to @c outb() in
 * @c interrupts.c but is a separate inline to avoid a cross-unit dependency
 * in the UEFI loader compilation unit.
 *
 * Only compiled when @c ARCH_AMD64 is defined.
 *
 * @param port 16-bit I/O port address (e.g., @c COM1_BASE + 0 = 0x3F8).
 * @param val  Byte value to write.
 */
static inline void raw_outb(uint16_t port, uint8_t val)
{
    __asm__ volatile ("outb %0, %1" : : "a"(val), "Nd"(port));
}

/**
 * @brief Read a byte from an x86 I/O port via @c IN (early-boot raw path).
 *
 * Used by @c raw_serial_putc() to poll the UART Line Status Register
 * (COM1 + 5) for the Transmitter Holding Register Empty (THRE) bit
 * before writing a character, preventing serial output corruption.
 * Only compiled when @c ARCH_AMD64 is defined.
 *
 * @param port 16-bit I/O port address (e.g., @c COM1_BASE + 5 = 0x3FD).
 * @return Byte value read from the port.
 */
static inline uint8_t raw_inb(uint16_t port)
{
    uint8_t ret;
    __asm__ volatile ("inb %1, %0" : "=a"(ret) : "Nd"(port));
    return ret;
}

/**
 * @brief Initialise COM1 (0x3F8) to 115200 8N1 before UEFI console exits.
 *
 * Programs the 16550-compatible UART at base address @c COM1_BASE using
 * direct port I/O. Sequence:
 * 1. Disable all UART interrupts (IER = 0x00).
 * 2. Assert DLAB (Divisor Latch Access Bit) to access the divisor registers.
 * 3. Write divisor = 1 (115200 baud at 115200 Hz base clock) to DLL and DLH.
 * 4. Clear DLAB; configure 8N1 (8 data bits, no parity, 1 stop bit; LCR=0x03).
 * 5. Enable FIFO, clear TX/RX FIFOs, set 14-byte interrupt threshold (FCR=0xC7).
 * 6. Assert RTS and DSR (MCR=0x0B).
 *
 * Called once at the beginning of @c efi_main() on amd64 to ensure the
 * @c RAW_LOG() macro can emit diagnostic messages throughout the boot
 * sequence, including phases where the UEFI console is no longer available.
 * Only compiled when @c ARCH_AMD64 is defined.
 */
static void raw_serial_init(void)
{
    /* Disable interrupts */
    raw_outb(COM1_BASE + 1, 0x00);
    /* Enable DLAB */
    raw_outb(COM1_BASE + 3, 0x80);
    /* Divisor 1 = 115200 baud */
    raw_outb(COM1_BASE + 0, 0x01);
    raw_outb(COM1_BASE + 1, 0x00);
    /* 8N1 */
    raw_outb(COM1_BASE + 3, 0x03);
    /* Enable FIFO, clear, 14-byte threshold */
    raw_outb(COM1_BASE + 2, 0xC7);
    /* RTS/DSR set */
    raw_outb(COM1_BASE + 4, 0x0B);
}

/**
 * @brief Write a single character to COM1 with THRE polling.
 *
 * Spins on bit 5 (Transmitter Holding Register Empty) of the Line Status
 * Register (COM1 + 5) until the UART is ready to accept a new byte, then
 * writes @p c to the Transmitter Holding Register (COM1 + 0). This busy-
 * wait is acceptable in the UEFI loader phase because interrupts are
 * either managed by UEFI or not yet configured by the kernel.
 *
 * Only compiled when @c ARCH_AMD64 is defined.
 *
 * @param c Character to transmit.
 */
static void raw_serial_putc(char c)
{
    while ((raw_inb(COM1_BASE + 5) & 0x20) == 0) { }
    raw_outb(COM1_BASE + 0, (uint8_t)c);
}

/**
 * @brief Write a NUL-terminated string to COM1 with implicit LF→CRLF conversion.
 *
 * Iterates over @p s and calls @c raw_serial_putc() for each character.
 * A bare @c '\\n' is preceded by a @c '\\r' to produce proper CRLF line
 * endings expected by serial terminals. Used via the @c RAW_LOG() macro
 * for diagnostic output before the kernel console subsystem is live.
 *
 * Only compiled when @c ARCH_AMD64 is defined.
 *
 * @param s NUL-terminated string to transmit.
 */
static void raw_serial_puts(const char *s)
{
    while (*s)
    {
        char c = *s++;
        if (c == '\n') raw_serial_putc('\r');
        raw_serial_putc(c);
    }
}
#define RAW_LOG(str) raw_serial_puts(str)

#elif defined(__riscv) || defined(ARCH_RISCV64)

/*
 * QEMU's riscv "virt" machine exposes a 16550-compatible UART (OpenSBI
 * reports "Platform Console Device: uart8250") as byte-addressed MMIO at
 * 0x10000000 (matches Domain0 Region03 in the OpenSBI boot banner). This
 * writes directly to the UART registers, exactly like the amd64
 * raw_serial_* helpers above, so RAW_LOG() actually produces output on
 * riscv64 instead of silently no-op'ing (previously the case for every
 * arch except amd64).
 */
#define UART_MMIO_BASE 0x10000000UL

static inline void raw_mmio_outb(uint64_t addr, uint8_t val)
{
    *(volatile uint8_t *)addr = val;
}

static inline uint8_t raw_mmio_inb(uint64_t addr)
{
    return *(volatile uint8_t *)addr;
}

static void raw_serial_putc(char c)
{
    while ((raw_mmio_inb(UART_MMIO_BASE + 5) & 0x20) == 0) { }
    raw_mmio_outb(UART_MMIO_BASE + 0, (uint8_t)c);
}

static void raw_serial_puts(const char *s)
{
    while (*s)
    {
        char c = *s++;
        if (c == '\n') raw_serial_putc('\r');
        raw_serial_putc(c);
    }
}
#define RAW_LOG(str) raw_serial_puts(str)

#else
#define RAW_LOG(str) ((void)0)
#endif

static BootInfo g_boot_info = {0};

/*
 * Numbered UEFI debug checkpoints (pre-ExitBootServices bring-up aid).
 *
 * First-boot debugging on hardware with no serial/UART and no framebuffer
 * driver yet: ConOut (the UEFI firmware's own text console on the
 * HDMI/DP-connected monitor) is the only output available before
 * ExitBootServices(). Each checkpoint prints "[CKPT nnn] <label>" and
 * stalls a couple of seconds so a photo can be taken, then the boot
 * continues. Purely additive instrumentation for a single debugging pass;
 * set UEFI_DEBUG_CHECKPOINTS to 0 (or pass -DUEFI_DEBUG_CHECKPOINTS=0) to
 * strip it back out once real console/logging exists.
 *
 * Deliberately confined to the region before the final GetMemoryMap() /
 * ExitBootServices() pair below: that region must call nothing that could
 * perturb the memory map (see the "DO NOTHING" comment on Phase B), so no
 * checkpoint fires between GetMemoryMap() and ExitBootServices(), and none
 * fires after ExitBootServices() returns (ConOut is not guaranteed usable
 * post-EBS on real firmware, and this loader has no serial fallback that
 * would be visible on the target board anyway).
 *
 * Checkpoint list (see also the cross-reference note in this file's
 * directory README / commit message):
 *   1. Entered efi_main - ConOut live
 *   2. Serial (COM1) initialized                 [ARCH_AMD64 only]
 *   3. kernel.elf loaded from ESP - OK            [!MONOLITHIC_BUILD only]
 *   4. Command line parsed
 *   5. Kernel stack allocation decided
 *   6. Boot info collected (ACPI table located)
 *   7. GOP query complete (outcome-specific label)
 *   8. About to enter ExitBootServices retry loop
 */
#ifndef UEFI_DEBUG_CHECKPOINTS
#define UEFI_DEBUG_CHECKPOINTS 1
#endif

#if UEFI_DEBUG_CHECKPOINTS
#define DEBUG_CHECKPOINT_STALL_US (2u * 1000u * 1000u)

static void debug_checkpoint(EFI_SYSTEM_TABLE *SystemTable, UINT32 id, const CHAR16 *label)
{
    CHAR16 line[192];
    UINTN pos = 0;
    UINTN i;
    UINTN val;

    static const CHAR16 prefix[] = L"[CKPT ";
    static const CHAR16 mid[] = L"] ";
    static const CHAR16 nl[] = L"\r\n";

    for (i = 0; prefix[i] != 0; ++i) line[pos++] = prefix[i];

    /* zero-padded 3-digit checkpoint number */
    val = id;
    {
        CHAR16 digits[3];
        int d;
        for (d = 2; d >= 0; --d) {
            digits[d] = (CHAR16)(L'0' + (val % 10));
            val /= 10;
        }
        for (d = 0; d < 3; ++d) line[pos++] = digits[d];
    }

    for (i = 0; mid[i] != 0; ++i) line[pos++] = mid[i];

    if (label != NULL) {
        for (i = 0; label[i] != 0 && pos < (sizeof(line) / sizeof(line[0])) - 3; ++i) {
            line[pos++] = label[i];
        }
    }

    for (i = 0; nl[i] != 0; ++i) line[pos++] = nl[i];
    line[pos] = 0;

    if (SystemTable != NULL && SystemTable->ConOut != NULL) {
        SystemTable->ConOut->OutputString(SystemTable->ConOut, line);
    }

    if (SystemTable != NULL && SystemTable->BootServices != NULL &&
        SystemTable->BootServices->Stall != NULL) {
        SystemTable->BootServices->Stall(DEBUG_CHECKPOINT_STALL_US);
    }
}
#else
#define debug_checkpoint(SystemTable, id, label) ((void)0)
#endif /* UEFI_DEBUG_CHECKPOINTS */

/**
 * @brief Compare two EFI GUIDs for equality.
 *
 * Compares all 128 bits of the two GUIDs field by field:
 * @c Data1 (32-bit), @c Data2 (16-bit), @c Data3 (16-bit), and all
 * 8 bytes of @c Data4. Returns non-zero only if every field matches.
 *
 * Used by @c efi_main() to search the UEFI @c ConfigurationTable for
 * the ACPI 2.0 and ACPI 1.0 table GUIDs (@c EFI_ACPI_20_TABLE_GUID and
 * @c EFI_ACPI_TABLE_GUID) to populate @c g_boot_info.acpi_table before
 * calling @c ExitBootServices().
 *
 * @param a Pointer to the first EFI GUID.
 * @param b Pointer to the second EFI GUID.
 * @return Non-zero if the GUIDs are equal, 0 otherwise.
 */
static int guid_equals(const EFI_GUID* a, const EFI_GUID* b)
{
    return a->Data1 == b->Data1 &&
           a->Data2 == b->Data2 &&
           a->Data3 == b->Data3 &&
           a->Data4[0] == b->Data4[0] &&
           a->Data4[1] == b->Data4[1] &&
           a->Data4[2] == b->Data4[2] &&
           a->Data4[3] == b->Data4[3] &&
           a->Data4[4] == b->Data4[4] &&
           a->Data4[5] == b->Data4[5] &&
           a->Data4[6] == b->Data4[6] &&
           a->Data4[7] == b->Data4[7];
}

/**
 * @brief Read @c /starforth.cfg from the EFI System Partition.
 *
 * Opens the ESP root directory using @c EFI_SIMPLE_FILE_SYSTEM_PROTOCOL
 * via the device handle of the loaded image, then reads the file
 * @c "starforth.cfg" into @p buf. On success, @p buf contains a
 * NUL-terminated ASCII command-line string (at most @p buflen - 1 bytes)
 * with trailing CR and LF characters stripped.
 *
 * This is the second-priority command-line source in @c efi_main() (after
 * the NVRAM one-shot variable @c StarForthBootArgs). It is intended for
 * installer-supplied or build-time-configured boot options that persist
 * across reboots, unlike the ephemeral NVRAM variable.
 *
 * Both @c MONOLITHIC_BUILD and split-build paths call this function;
 * it is not gated by a preprocessor condition.
 *
 * @param ImageHandle UEFI image handle for the current application;
 *                    used to retrieve @c EFI_LOADED_IMAGE_PROTOCOL.
 * @param BS          UEFI Boot Services table pointer; must not be NULL.
 * @param buf         Caller-supplied buffer; receives the NUL-terminated
 *                    ASCII config string on success.
 * @param buflen      Size of @p buf in bytes; at most @p buflen-1 bytes
 *                    of the file are read.
 * @return 1 if the file was read and contained non-empty content;
 *         0 on any UEFI error or if the file is empty after stripping.
 */
static int load_cfg_from_esp(EFI_HANDLE ImageHandle, EFI_BOOT_SERVICES *BS,
                              char *buf, UINTN buflen)
{
    EFI_LOADED_IMAGE_PROTOCOL        *loaded_image = NULL;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL  *fs           = NULL;
    EFI_FILE_PROTOCOL                *root         = NULL;
    EFI_FILE_PROTOCOL                *cfg_file     = NULL;
    EFI_STATUS status;
    UINTN read_size;

    buf[0] = '\0';

    status = BS->HandleProtocol(ImageHandle,
                                (EFI_GUID *)&EFI_LOADED_IMAGE_PROTOCOL_GUID,
                                (void **)&loaded_image);
    if (status != EFI_SUCCESS || !loaded_image) return 0;

    status = BS->HandleProtocol(loaded_image->DeviceHandle,
                                (EFI_GUID *)&EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_GUID,
                                (void **)&fs);
    if (status != EFI_SUCCESS || !fs) return 0;

    status = fs->OpenVolume(fs, &root);
    if (status != EFI_SUCCESS || !root) return 0;

    status = root->Open(root, &cfg_file, L"starforth.cfg", EFI_FILE_MODE_READ, 0);
    if (status != EFI_SUCCESS) {
        root->Close(root);
        return 0;
    }

    read_size = buflen - 1;
    status = cfg_file->Read(cfg_file, &read_size, buf);
    cfg_file->Close(cfg_file);
    root->Close(root);

    if (status != EFI_SUCCESS) {
        buf[0] = '\0';
        return 0;
    }

    buf[read_size] = '\0';
    /* Strip trailing CR/LF */
    while (read_size > 0 &&
           (buf[read_size - 1] == '\n' || buf[read_size - 1] == '\r'))
        buf[--read_size] = '\0';

    return read_size > 0;
}

#ifndef MONOLITHIC_BUILD
/* Kernel entry point signature */
typedef void (*KernelEntry)(BootInfo *boot_info);

/* Kernel ELF buffer - dynamically allocated via UEFI Boot Services */
#define KERNEL_MAX_SIZE (8 * 1024 * 1024)
static uint8_t *kernel_elf_buffer = NULL;
static uint64_t kernel_elf_size = 0;

/**
 * @brief Load @c kernel.elf from the ESP into UEFI-allocated memory (split build).
 *
 * Used only in the split-build path (@c !MONOLITHIC_BUILD) where the kernel
 * is a separate ELF file on the EFI System Partition rather than being
 * linked directly into the loader image. Execution order:
 *
 * 1. Allocates @c KERNEL_MAX_SIZE (8 MB) of @c EfiLoaderData pages via
 *    @c AllocatePages. Stores the result in the module-static
 *    @c kernel_elf_buffer pointer.
 * 2. Retrieves @c EFI_LOADED_IMAGE_PROTOCOL from @p ImageHandle to obtain
 *    the boot device handle.
 * 3. Opens @c EFI_SIMPLE_FILE_SYSTEM_PROTOCOL on the device handle and
 *    opens the ESP volume root directory.
 * 4. Opens @c "kernel.elf" for read-only access.
 * 5. Queries file size via @c EFI_FILE_INFO; returns
 *    @c EFI_BUFFER_TOO_SMALL if the file exceeds @c KERNEL_MAX_SIZE.
 * 6. Reads the entire file into @c kernel_elf_buffer and stores the byte
 *    count in @c kernel_elf_size.
 *
 * Must be called before @c ExitBootServices() — file system access via
 * Boot Services is unavailable after EBS. The allocated buffer and size
 * are then passed to @c elf_load_kernel() in Phase C.
 *
 * Only compiled when @c MONOLITHIC_BUILD is not defined.
 *
 * @param ImageHandle UEFI image handle for the current application.
 * @param BS          UEFI Boot Services table pointer.
 * @return @c EFI_SUCCESS on success, or a UEFI status code on failure.
 */
static EFI_STATUS load_kernel_from_esp(EFI_HANDLE ImageHandle, EFI_BOOT_SERVICES *BS)
{
    EFI_STATUS status;
    EFI_LOADED_IMAGE_PROTOCOL *loaded_image = NULL;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *fs = NULL;
    EFI_FILE_PROTOCOL *root = NULL;
    EFI_FILE_PROTOCOL *kernel_file = NULL;

    RAW_LOG("Loading kernel.elf from ESP...\n");

    /* Allocate buffer for kernel ELF using UEFI Boot Services */
    UINTN pages_needed = (KERNEL_MAX_SIZE + 4095) / 4096;
    EFI_PHYSICAL_ADDRESS buffer_addr = 0;
    status = BS->AllocatePages(AllocateAnyPages, EfiLoaderData, pages_needed, &buffer_addr);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to allocate kernel buffer\n");
        return status;
    }
    kernel_elf_buffer = (uint8_t *)(uintptr_t)buffer_addr;
    RAW_LOG("Kernel buffer allocated\n");

    /* Get loaded image protocol */
    status = BS->HandleProtocol(ImageHandle, (EFI_GUID *)&EFI_LOADED_IMAGE_PROTOCOL_GUID,
                                (void **)&loaded_image);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to get LoadedImageProtocol\n");
        return status;
    }

    /* Get file system protocol from device handle */
    status = BS->HandleProtocol(loaded_image->DeviceHandle,
                                (EFI_GUID *)&EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_GUID,
                                (void **)&fs);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to get FileSystemProtocol\n");
        return status;
    }

    /* Open volume (root directory) */
    status = fs->OpenVolume(fs, &root);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to open volume\n");
        return status;
    }

    /* Open kernel.elf file */
    status = root->Open(root, &kernel_file, L"kernel.elf",
                       EFI_FILE_MODE_READ, 0);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to open kernel.elf\n");
        root->Close(root);
        return status;
    }

    /* Get file size */
    EFI_FILE_INFO file_info_buffer[128]; /* Should be enough for file info */
    UINTN file_info_size = sizeof(file_info_buffer);
    status = kernel_file->GetInfo(kernel_file, (EFI_GUID *)&EFI_FILE_INFO_GUID,
                                  &file_info_size, file_info_buffer);
    if (status != EFI_SUCCESS) {
        RAW_LOG("Failed to get file info\n");
        kernel_file->Close(kernel_file);
        root->Close(root);
        return status;
    }

    EFI_FILE_INFO *file_info = (EFI_FILE_INFO *)file_info_buffer;
    kernel_elf_size = file_info->FileSize;

    /* Check if kernel fits in buffer */
    if (kernel_elf_size > KERNEL_MAX_SIZE) {
        RAW_LOG("Kernel too large for buffer\n");
        kernel_file->Close(kernel_file);
        root->Close(root);
        return EFI_BUFFER_TOO_SMALL;
    }

    /* Read kernel into buffer */
    UINTN read_size = kernel_elf_size;
    status = kernel_file->Read(kernel_file, &read_size, kernel_elf_buffer);
    if (status != EFI_SUCCESS || read_size != kernel_elf_size) {
        RAW_LOG("Failed to read kernel.elf\n");
        kernel_file->Close(kernel_file);
        root->Close(root);
        return status != EFI_SUCCESS ? status : EFI_ABORTED;
    }

    RAW_LOG("Kernel loaded successfully\n");

    /* Close files */
    kernel_file->Close(kernel_file);
    root->Close(root);

    return EFI_SUCCESS;
}
#endif /* !MONOLITHIC_BUILD */

/**
 * @brief UEFI application entry point — boot loader for StarKernel.
 *
 * @c efi_main() is the UEFI PE32+ entry point called by the firmware
 * immediately after the image is loaded and validated. It drives the full
 * boot sequence in two phases:
 *
 * **Phase A — UEFI Boot Services are active:**
 *
 * A-0. Resets the UEFI ConOut console and emits an identification banner.
 *      On amd64, initialises COM1 via @c raw_serial_init() for raw-serial
 *      boot logging (@c RAW_LOG() macro).
 *
 * A-1. Assembles the ASCII kernel command line from three sources in
 *      priority order:
 *      1. @c StarForthBootArgs NVRAM variable (one-shot, set by REBOOT word;
 *         immediately cleared after reading to prevent re-use on next boot).
 *      2. @c /starforth.cfg on the ESP (installer / build-time config).
 *      3. @c LoadOptions from the UEFI boot manager (typically the EFI shell
 *         command line or boot entry options).
 *      Parses the result via @c cmdline_parse_ascii() into
 *      @c g_boot_info.args.
 *
 * A-2. If @c --stack=@<size@> was given, allocates @p stack_size bytes of
 *      @c EfiLoaderData pages for the kernel stack and stores the base
 *      address and size in @c g_boot_info.
 *
 * A-3. Pre-fills non-EBS-dependent @c BootInfo fields: runtime services
 *      pointer, ACPI table (found in @c ConfigurationTable by GUID),
 *      and GOP framebuffer geometry.
 *
 * A-4. (Split build only) Loads @c kernel.elf from the ESP via
 *      @c load_kernel_from_esp() into the pre-allocated 8 MB buffer.
 *
 * **Phase B — ExitBootServices:**
 *
 * Calls @c GetMemoryMap() to obtain the current map key, then immediately
 * calls @c ExitBootServices(). Retries up to 16 times on
 * @c EFI_INVALID_PARAMETER (map key stale) as required by the UEFI spec.
 * No UEFI Boot Services calls are made between @c GetMemoryMap and
 * @c ExitBootServices — this is the UEFI "golden path" requirement.
 *
 * **Phase C — Post-ExitBootServices:**
 *
 * - **Monolithic build**: calls @c kernel_main(@c &g_boot_info) directly
 *   (kernel is linked into the same PE image).
 * - **Split build**: parses and loads the kernel ELF via
 *   @c elf_load_kernel(), then performs an indirect call to the resolved
 *   @c kernel_main() entry point address.
 *
 * If the kernel returns (which it must not), @c efi_main() halts the
 * processor in an infinite @c arch_halt() loop.
 *
 * @param ImageHandle  UEFI handle for this loaded application image.
 * @param SystemTable  Pointer to the UEFI system table (Boot/Runtime
 *                     Services, ConOut, ConfigurationTable, etc.).
 * @return @c EFI_SUCCESS if the sequence completes without a fatal error
 *         before @c ExitBootServices (in practice, control never returns
 *         here because the kernel takes over permanently).
 */
EFI_STATUS EFIAPI efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS status;
    EFI_BOOT_SERVICES *BS = SystemTable->BootServices;

    /* Static map buffer (256 KiB) to avoid allocations in Phase B */
    static EFI_MEMORY_DESCRIPTOR memory_map_static[256 * 1024 / sizeof(EFI_MEMORY_DESCRIPTOR)];
    EFI_MEMORY_DESCRIPTOR *memory_map = memory_map_static;
    UINTN memory_map_capacity = sizeof(memory_map_static);

    UINTN map_key = 0;
    UINTN descriptor_size = 0;
    UINT32 descriptor_version = 0;
    int exited_boot_services = 0;

    /* Phase A: safe to use UEFI console and file services */
    SystemTable->ConOut->Reset(SystemTable->ConOut, FALSE);
    SystemTable->ConOut->OutputString(SystemTable->ConOut, L"StarKernel UEFI Loader\r\n");
    SystemTable->ConOut->OutputString(SystemTable->ConOut, L"Loading kernel from ESP...\r\n");
    debug_checkpoint(SystemTable, 1, L"Entered efi_main - ConOut live");

#if defined(ARCH_AMD64)
    raw_serial_init();
    RAW_LOG("RAW SERIAL UP\n");
    debug_checkpoint(SystemTable, 2, L"Serial (COM1) initialized");
#endif

#ifndef MONOLITHIC_BUILD
    /* Load kernel.elf from ESP BEFORE ExitBootServices */
    status = load_kernel_from_esp(ImageHandle, BS);
    if (status != EFI_SUCCESS) {
        SystemTable->ConOut->OutputString(SystemTable->ConOut, L"FATAL: Failed to load kernel.elf\r\n");
        RAW_LOG("FATAL: Failed to load kernel.elf\n");
        while (1) arch_halt();
        return status;
    }
    debug_checkpoint(SystemTable, 3, L"kernel.elf loaded from ESP - OK");
#else
    RAW_LOG("Monolithic build - kernel linked directly\n");
#endif

    SystemTable->ConOut->OutputString(SystemTable->ConOut, L"Collecting boot information...\r\n");

    /* ------------------------------------------------------------------ */
    /* Phase A-1: Parse kernel command line                                */
    /*                                                                     */
    /* Priority (highest first):                                           */
    /*   1. StarForthBootArgs NVRAM variable (one-shot, set by REBOOT)    */
    /*   2. EFI_LOADED_IMAGE_PROTOCOL->LoadOptions (firmware/boot mgr)    */
    /*   3. Empty string → all KernelArgs defaults                        */
    /* ------------------------------------------------------------------ */
    {
        char cmdline_buf[KERNEL_ARGS_CMDLINE_MAX];
        int  used_nvram = 0;
        int  used_cfg   = 0;

        cmdline_buf[0] = '\0';

        /* Priority 1: one-shot NVRAM variable (set by REBOOT word) */
        {
            EFI_GUID vendor_guid = STARFORTH_VENDOR_GUID;
            EFI_GET_VARIABLE GetVariable =
                (EFI_GET_VARIABLE)SystemTable->RuntimeServices->GetVariable;
            UINTN data_size = (UINTN)(KERNEL_ARGS_CMDLINE_MAX - 1);
            UINT32 attrs = 0;

            EFI_STATUS vs = GetVariable(
                (CHAR16 *)SF_VAR_BOOT_ARGS,
                &vendor_guid,
                &attrs,
                &data_size,
                cmdline_buf);

            if (vs == EFI_SUCCESS && data_size > 0) {
                cmdline_buf[data_size] = '\0';
                used_nvram = 1;
                RAW_LOG("CmdLine: from NVRAM StarForthBootArgs\n");

                /* One-shot: clear the variable immediately */
                EFI_SET_VARIABLE SetVariable =
                    (EFI_SET_VARIABLE)SystemTable->RuntimeServices->SetVariable;
                SetVariable(
                    (CHAR16 *)SF_VAR_BOOT_ARGS,
                    &vendor_guid,
                    EFI_VARIABLE_NON_VOLATILE |
                    EFI_VARIABLE_BOOTSERVICE_ACCESS |
                    EFI_VARIABLE_RUNTIME_ACCESS,
                    0, NULL);
            }
        }

        /* Priority 2: /starforth.cfg on the ESP (build-time / installer config) */
        if (!used_nvram) {
            if (load_cfg_from_esp(ImageHandle, BS,
                                  cmdline_buf, sizeof(cmdline_buf))) {
                used_cfg = 1;
                RAW_LOG("CmdLine: from starforth.cfg\n");
            }
        }

        /* Priority 3: LoadOptions from firmware / boot manager */
        if (!used_nvram && !used_cfg) {
            EFI_LOADED_IMAGE_PROTOCOL *loaded_image = NULL;
            EFI_STATUS li_status = BS->HandleProtocol(
                ImageHandle,
                (EFI_GUID *)&EFI_LOADED_IMAGE_PROTOCOL_GUID,
                (void **)&loaded_image);

            if (li_status == EFI_SUCCESS && loaded_image &&
                loaded_image->LoadOptionsSize > 0 && loaded_image->LoadOptions) {
                CHAR16 *opts = (CHAR16 *)loaded_image->LoadOptions;
                UINTN   n    = loaded_image->LoadOptionsSize / sizeof(CHAR16);
                UINTN   i;
                for (i = 0; i + 1 < (UINTN)(KERNEL_ARGS_CMDLINE_MAX) && i < n && opts[i]; i++)
                    cmdline_buf[i] = (char)(opts[i] & 0xFFu);
                cmdline_buf[i] = '\0';
                RAW_LOG("CmdLine: from LoadOptions\n");
            }
        }

        cmdline_parse_ascii(cmdline_buf, &g_boot_info.args);
        RAW_LOG("CmdLine: parsed OK\n");
    }
    debug_checkpoint(SystemTable, 4, L"Command line parsed");

    /* ------------------------------------------------------------------ */
    /* Phase A-2: Allocate dynamic kernel stack if --stack= was given     */
    /* ------------------------------------------------------------------ */
    {
        uint64_t stack_sz = g_boot_info.args.stack_size
                          ? g_boot_info.args.stack_size
                          : (uint64_t)0;

        g_boot_info.kernel_stack_base = NULL;
        g_boot_info.kernel_stack_size = 0;

        if (stack_sz > 0) {
            UINTN pages = (UINTN)((stack_sz + 4095u) / 4096u);
            EFI_PHYSICAL_ADDRESS stack_addr = 0;
            EFI_STATUS sa = BS->AllocatePages(
                AllocateAnyPages, EfiLoaderData, pages, &stack_addr);
            if (sa == EFI_SUCCESS) {
                g_boot_info.kernel_stack_base = (void *)(uintptr_t)stack_addr;
                g_boot_info.kernel_stack_size = (uint64_t)(pages * 4096u);
                RAW_LOG("Stack: dynamic stack allocated\n");
            } else {
                RAW_LOG("Stack: dynamic alloc failed — using BSS fallback\n");
            }
        }
    }
    debug_checkpoint(SystemTable, 5, L"Kernel stack allocation decided");

    /* Fill BootInfo fields that do NOT require EBS first */
    g_boot_info.runtime_services = SystemTable->RuntimeServices;
    g_boot_info.acpi_table = NULL;
    g_boot_info.framebuffer.base = NULL;
    g_boot_info.framebuffer.size = 0;
    g_boot_info.framebuffer.width = 0;
    g_boot_info.framebuffer.height = 0;
    g_boot_info.framebuffer.pixels_per_scanline = 0;
    g_boot_info.framebuffer.pixel_format = (UINT32)PixelBltOnly;
    g_boot_info.uefi_boot_services_exited = 0;

    /* Locate ACPI table (safe pre-EBS). Prefer the ACPI 2.0+ RSDP (has an
     * XSDT) over the legacy ACPI 1.0 one (RSDT only, revision 0) -- a
     * single OR'd loop that takes whichever GUID appears first in the
     * firmware's configuration table can silently hand back the 1.0
     * pointer even when a 2.0 one is also present, which then fails any
     * XSDT-based table lookup (e.g. FADT PM_TMR_BLK discovery in
     * timer_init()). Two explicit passes: 2.0 first, 1.0 only as fallback. */
    {
        EFI_CONFIGURATION_TABLE *config_tables =
            (EFI_CONFIGURATION_TABLE *)SystemTable->ConfigurationTable;
        void *acpi10_table = NULL;

        for (UINTN i = 0; i < SystemTable->NumberOfTableEntries; ++i) {
            if (guid_equals(&config_tables[i].VendorGuid, &EFI_ACPI_20_TABLE_GUID)) {
                g_boot_info.acpi_table = config_tables[i].VendorTable;
                break;
            }
            if (!acpi10_table && guid_equals(&config_tables[i].VendorGuid, &EFI_ACPI_TABLE_GUID)) {
                acpi10_table = config_tables[i].VendorTable;
            }
        }
        if (!g_boot_info.acpi_table) {
            g_boot_info.acpi_table = acpi10_table;
        }
    }
    debug_checkpoint(SystemTable, 6, L"Boot info collected (ACPI table located)");

    /* Locate GOP and populate framebuffer info (safe pre-EBS) */
    {
        EFI_GUID gop_guid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
        EFI_GRAPHICS_OUTPUT_PROTOCOL *gop = NULL;
        EFI_STATUS gop_status = BS->LocateProtocol(&gop_guid, NULL, (void **)&gop);
        if (gop_status == EFI_SUCCESS && gop && gop->Mode && gop->Mode->Info) {
            EFI_GRAPHICS_OUTPUT_PROTOCOL_MODE *mode = gop->Mode;
            if (mode->Info->PixelFormat != PixelBltOnly) {
                g_boot_info.framebuffer.base             = (void *)(UINTN)mode->FrameBufferBase;
                g_boot_info.framebuffer.size             = (UINTN)mode->FrameBufferSize;
                g_boot_info.framebuffer.width            = mode->Info->HorizontalResolution;
                g_boot_info.framebuffer.height           = mode->Info->VerticalResolution;
                g_boot_info.framebuffer.pixels_per_scanline = mode->Info->PixelsPerScanLine;
                g_boot_info.framebuffer.pixel_format     = (UINT32)mode->Info->PixelFormat;
                RAW_LOG("GOP: linear framebuffer found\n");
                debug_checkpoint(SystemTable, 7, L"GOP: linear framebuffer found");
            } else {
                RAW_LOG("GOP: PixelBltOnly - no linear framebuffer\n");
                debug_checkpoint(SystemTable, 7, L"GOP: BltOnly, no linear fb");
            }
        } else {
            RAW_LOG("GOP: protocol not found\n");
            debug_checkpoint(SystemTable, 7, L"GOP: protocol not found");
        }
    }

    /*
     * Phase B: ExitBootServices loop
     * RULE: Between the final GetMemoryMap and ExitBootServices => DO NOTHING.
     * No ConOut, no extra services, no protocol opens, no "verification" GetMemoryMap.
     *
     * (Checkpoint 8 fires immediately below, before this loop starts, rather
     * than between GetMemoryMap and ExitBootServices, specifically to respect
     * this rule -- see the debug_checkpoint doc comment near the top of this
     * file.)
     */
    debug_checkpoint(SystemTable, 8, L"About to enter ExitBootServices retry loop");
    for (int attempt = 0; attempt < 16; ++attempt) {
        UINTN map_sz = memory_map_capacity;

        status = BS->GetMemoryMap(
            &map_sz,
            memory_map,
            &map_key,
            &descriptor_size,
            &descriptor_version
        );

        if (status == EFI_BUFFER_TOO_SMALL) {
            RAW_LOG("GetMemoryMap: BUFFER_TOO_SMALL\r\n");
            return status;
        }

        if (status != EFI_SUCCESS) {
            RAW_LOG("GetMemoryMap: ERROR\r\n");
            return status;
        }

        /* Stash map into BootInfo (still pre-EBS) */
        g_boot_info.memory_map = memory_map;
        g_boot_info.memory_map_size = map_sz;
        g_boot_info.memory_map_descriptor_size = descriptor_size;
        g_boot_info.runtime_services = SystemTable->RuntimeServices;

#if defined(ARCH_AMD64)
        RAW_LOG("EBS...\r\n");
#endif

        /* Call ExitBootServices immediately after GetMemoryMap */
        status = BS->ExitBootServices(ImageHandle, map_key);
        if (status == EFI_SUCCESS) {
            exited_boot_services = 1;
#if defined(ARCH_AMD64)
            RAW_LOG("EBS OK\r\n");
#endif
            break;
        }

        if (status != EFI_INVALID_PARAMETER) {
#if defined(ARCH_AMD64)
            RAW_LOG("EBS fail non-invalid\r\n");
#endif
            return status;
        }

#if defined(ARCH_AMD64)
        RAW_LOG("EBS invalid -> retry\r\n");
#endif
        /* loop will retry */
    }

    g_boot_info.uefi_boot_services_exited = exited_boot_services ? 1u : 0u;

#ifdef MONOLITHIC_BUILD
    /*
     * Monolithic build: kernel is linked directly, call kernel_main
     */
    RAW_LOG("Calling kernel_main (monolithic)...\n");
    kernel_main(&g_boot_info);
#else
    /*
     * Phase C: Post-ExitBootServices - Parse and load kernel
     * BootServices are GONE, can only use runtime services
     */
    RAW_LOG("Parsing kernel ELF...\n");

    Elf64_Addr entry_point = 0;
    if (!elf_load_kernel(kernel_elf_buffer, kernel_elf_size, &entry_point)) {
        RAW_LOG("FATAL: Failed to load kernel ELF\n");
        while (1) arch_halt();
    }

    RAW_LOG("Jumping to kernel entry point...\n");

    /* Jump to kernel entry point */
    KernelEntry kernel_entry = (KernelEntry)entry_point;
    kernel_entry(&g_boot_info);
#endif

    /* Should never reach here */
    RAW_LOG("FATAL: Kernel returned\n");
    while (1) {
        arch_halt();
    }

    return EFI_SUCCESS;
}