/* =============================================================================
 * main.c — RNG UEFI bootloader (rsbl) — minimal smoke-test version
 *
 * Just prints to screen then loads kernel + initrd via LoadImage/StartImage.
 * ============================================================================= */

#include "uefi.h"

static const char16 kpath[]   = u"\\bzImage";
static const char16 ipath[]   = u"\\initramfs.cpio.gz";
/* same cmdline systemd-boot used */
static const char16 base_opt[] = u"console=ttyS0,115200 console=tty0 earlyprintk=serial,ttyS0,115200 loglevel=3 rng_tty=ttyS0 initrd=\\initramfs.cpio.gz";

static efi_guid lg = EFI_LOADED_IMAGE_GUID;
static efi_guid sg = EFI_SIMPLE_FILE_SYSTEM_GUID;

static void pr(efi_simple_text_output *c, const char16 *s) { c->output_string(c, s); }
static void ph(efi_simple_text_output *c, u64 v) {
    char16 b[19] = u"0x0000000000000000";
    for (int i = 17; i >= 2; i--) {
        u8 n = (u8)(v & 0xF);
        b[i] = n < 10 ? u'0' + n : u'a' + n - 10;
        v >>= 4;
    }
    c->output_string(c, b);
}

static u64 allocate_pages(efi_boot_services *b, u64 mt, u64 np, u64 *a) {
    typedef u64 (*F)(u64, u64, u64, u64 *);
    return ((F)b->allocate_pages)(0, mt, np, a);
}
static u8 *rf(efi_boot_services *b, efi_file_protocol *f, u64 *sz, u64 mx) {
    u64 np = (mx + 4095) / 4096, a;
    /* LoaderData (2): scratch buffers for the OS loader, reclaimable by the
     * OS once it takes control.  Type 6 (RuntimeData) would mark them as
     * runtime-services memory that must survive forever. */
    if (allocate_pages(b, 2, np, &a)) return 0;
    u8 *d = (u8 *)a;
    u64 r = mx;
    if (f->read(f, &r, d) && r != mx) return 0;
    *sz = r; return d;
}

u64 __attribute__((aligned(16))) efi_main(handle ih, efi_system_table *st) {
    efi_simple_text_output *c = st->con_out;
    efi_boot_services      *b = st->boot_services;
    u64 s;

    pr(c, u"\r\n");
    pr(c, u"rsboot v0.1 -- RaiserSoft UEFI Bootloader\r\n");
    pr(c, u"\r\n");

    efi_loaded_image *our_li;
    s = b->handle_protocol(ih, &lg, (void **)&our_li);
    if (s) { pr(c, u"  handle_protocol LI failed\r\n"); goto h; }

    efi_simple_file_system *fs;
    s = b->handle_protocol(our_li->device_handle, &sg, (void **)&fs);
    if (s) { pr(c, u"  SFSP not on device handle\r\n"); goto h; }

    efi_file_protocol *rt;
    s = fs->open_volume(fs, (void **)&rt);
    if (s) { pr(c, u"  open_volume failed\r\n"); goto h; }

    pr(c, u"  [     loading kernel\r");
    efi_file_protocol *f;
    s = rt->open(rt, &f, (char16 *)kpath, 1, 0);
    if (s) { pr(c, u"  [FAIL] loading kernel\r\n"); goto h; }
    u64 ks; u8 *kd = rf(b, f, &ks, 64*1024*1024);
    f->close(f);
    if (!kd) { pr(c, u"  [FAIL] loading kernel\r\n"); goto h; }
    pr(c, u"  [  ok] loading kernel  "); ph(c, ks); pr(c, u" bytes\r\n");

    pr(c, u"  [     loading initrd\r");
    s = rt->open(rt, &f, (char16 *)ipath, 1, 0);
    if (s) { pr(c, u"  [FAIL] loading initrd\r\n"); goto h; }
    u64 is; u8 *id = rf(b, f, &is, 32*1024*1024);
    f->close(f);
    if (!id) { pr(c, u"  [FAIL] loading initrd\r\n"); goto h; }
    pr(c, u"  [  ok] loading initrd  "); ph(c, is); pr(c, u" bytes\r\n");

    pr(c, u"  [     starting image\r");
    handle kh = 0;
    s = b->load_image(0, ih, 0, kd, ks, &kh);
    if (s) { pr(c, u"  [FAIL] LoadImage\r\n"); goto h; }

    efi_loaded_image *kli;
    s = b->handle_protocol(kh, &lg, (void **)&kli);
    if (s) { pr(c, u"  [FAIL] kernel image info\r\n"); goto h; }
    kli->device_handle = our_li->device_handle;

    u16 osz = 0;
    while (base_opt[osz]) osz++;
    osz = (osz + 1) * 2;
    u16 *opts;
    {
        typedef u64 (*ap_t)(u8, u64, void **);
        s = ((ap_t)b->allocate_pool)(0, osz, (void **)&opts);
    }
    if (s) { pr(c, u"  [FAIL] cmdline alloc\r\n"); goto h; }
    for (u16 i = 0; i < osz / 2; i++) opts[i] = base_opt[i];
    kli->load_options = opts;
    kli->load_options_size = osz;

    pr(c, u"  [  ok] starting image\r\n");
    pr(c, u"\r\n");
    pr(c, u"  Launching the kernel ...\r\n");
    pr(c, u"\r\n");

    s = b->start_image(kh, 0, 0);
    pr(c, u"  kernel returned "); ph(c, s); pr(c, u"\r\n");

h:
    pr(c, u"  Halt\r\n");
    while (1) asm volatile("hlt");
    return s;
}