/*
 * Copyright (c) 2026 Netflix, Inc. Written by Warner Losh
 *
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Derived from memdisk_uefi.c
 * Copyright 2025 Richard Russo
 * SPDX-License-Identifier: BSD-2-Clause-Patent
 */

#include "loader_efi.h"
#include <efilib.h>
#include <Protocol/RamDisk.h>
#include "decompress.h"
#include <ipxe_download.h>
#include <sys/_param.h>

#define ULL(x) ((unsigned long long)(x))

static EFI_GUID ipxeGuid = IPXE_DOWNLOAD_PROTOCOL_GUID;
static EFI_GUID ramdiskGuid = EFI_RAM_DISK_PROTOCOL_GUID;
static EFI_GUID virtual_disk_guid = EFI_VIRTUAL_DISK_GUID;
static EFI_GUID virtual_cd_guid = EFI_VIRTUAL_CD_GUID;

static IPXE_DOWNLOAD_PROTOCOL *ipxe_download;
static EFI_RAM_DISK_PROTOCOL *ram_disk;

struct dl_state;
typedef struct dl_state dl_state;

static struct dl_state
{
        bool in_progress;
        size_t size;
        EFI_STATUS status;
        decomp_state *dctx;
} dl;

static void
download_cleanup(dl_state *ctx)
{
        if (ctx->dctx)
                decomp_fini(ctx->dctx, true);
        ctx->in_progress = false;
}

static EFI_STATUS EFIAPI
download_data(IN VOID *Context, IN VOID *Buffer, IN UINTN BufferLength, IN UINTN FileOffset)
{
        dl_state *ctx = Context;
        decomp_state *dctx = ctx->dctx;

        if (FileOffset == 0 && BufferLength == 0) {
                printf("Staritng the download\n");
                return (EFI_SUCCESS);
        }

        /*
         * Make a note of the size when we're hinted about it.
         */
        if (BufferLength == 0) {
                printf("We know we will download %llu bytes\n", ULL(FileOffset));
                ctx->size = FileOffset;
                ctx->status = EFI_SUCCESS;
                return (EFI_SUCCESS);
        }

        /*
         * Peek into the first chunk to see the format of the data.
         */
        if (FileOffset == 0) {
                dctx = decomp_init((uint8_t *)Buffer, (size_t)BufferLength, ctx->size);
                if (dctx == NULL) {
                        ctx->in_progress = false;
                        ctx->status = EFI_VOLUME_CORRUPTED;
                        return (ctx->status);
                }
                ctx->dctx = dctx;
        }

        enum step_return sr = decomp_step(dctx, Buffer, BufferLength, FileOffset);
        if (sr == err) {
                printf("Error on download\n");
                decomp_fini(dctx, true);
                return (EFI_VOLUME_CORRUPTED);
        }

        unsigned long long sofar = FileOffset + BufferLength;
#define MB  1000000
        if (sofar / MB != FileOffset / MB) {
                if (ctx->size)
                        printf("%dMB / %dMB (%d%%)\r",
                            (int)(sofar / MB),
                            (int)(ctx->size / MB),
                            (int)(100 * sofar / ctx->size));
                else
                        printf("%dMB\r", (int)(sofar / MB));
        }
        return (EFI_SUCCESS);
}

static void EFIAPI
download_finish(IN VOID *Context, IN EFI_STATUS Status)
{
        dl_state *ctx = Context;

        ctx->in_progress = false;
        ctx->status = Status;
        if (ctx->dctx)
                decomp_fini(ctx->dctx, EFI_ERROR(Status));
}

static void
do_download_ramdisk(CHAR8 *url, bool is_disk)
{
        EFI_STATUS Status;
        EFI_GUID disk_type = is_disk ? virtual_disk_guid : virtual_cd_guid;
        EFI_DEVICE_PATH_PROTOCOL *ram_disk_path;
        IPXE_DOWNLOAD_FILE token;
        dl_state *ctx = &dl;

        Status = BS->LocateProtocol(&ipxeGuid, NULL, (void**)&ipxe_download);
        if (EFI_ERROR(Status))
                return; /* most uses won't have this, don't whine */
        Status = BS->LocateProtocol(&ramdiskGuid, NULL, (void**)&ram_disk);
        if (EFI_ERROR(Status))
                return; /* XXX whine about it? */

        printf("Downloading %s as a %s\n", url, is_disk ? "disk" : "cd");
        ctx->in_progress = true;
        Status = ipxe_download->Start(ipxe_download, url, download_data, download_finish,
            &dl, &token);
        if (EFI_ERROR(Status)) {
                printf("Couldn't start download %u\n", (unsigned)Status);
                download_cleanup(ctx);
                return;
        }
        while (ctx->in_progress) {
                ipxe_download->Poll(ipxe_download);
        }
        if (EFI_ERROR(ctx->status)) {
                printf("Download had error %u\n", (unsigned)ctx->status);
                download_cleanup(ctx);
                return;
        }
        if (ctx->size == 0) {
                printf("Nothing downloaded\n");
                download_cleanup(ctx);
                return;
        }

        printf("\nDownloaded %llu bytes, actual size %llu -- registering ramdisk\n",
            ULL(ctx->size), ULL(decomp_buffer_length(ctx->dctx)));

        /*
         * Register the RamDisk with UEFI. This registers it so the rest of the
         * boot loader can see it as a block device.
         */
        Status = ram_disk->Register(decomp_buffer(ctx->dctx), decomp_buffer_length(ctx->dctx),
            &disk_type, NULL, &ram_disk_path);
        if (EFI_ERROR(Status)) {
                printf("failed to register ram disk %u\n", (unsigned)Status);
                download_cleanup(ctx);
                return;
        }

        CHAR16 *text = efi_devpath_name(ram_disk_path);
        if (text != NULL) {
                CHAR8 uefi_path[1024];
                printf("Installed RAM disk as %S\n", text);

                cpy16to8(text, uefi_path, sizeof(uefi_path));
                setenv("uefi_ignore_boot_mgr", "true", 1);
                setenv("uefi_rootdev", uefi_path, 1);
                efi_free_devpath_name(text);
        } else {
                printf("Installed RAM disk to unknown device type\n");
        }
}

/*
 * Scan the command line for memdisk=url or memcd=url. Do nothing if that's not
 * present, otherwise try to download that image.
 *
 * Open Question: Do we want some way to chain boot into the /boot/loader.efi or
 * \efi\boot\bootXXXXX.efi inside the ram disk we load? If so, how do we keep
 * from infinite chainbooting? Also, I don't understand the load it but don't save
 * it option...
 */
void
maybe_download_ramdisk(int argc, CHAR16 **argv)
{
        char var[256];

        for (int i = 0; i < argc; i++) {
                cpy16to8(argv[i], var, sizeof(var));
                if (strncmp(var, "memdisk=", 8) == 0) {
                        do_download_ramdisk(var + 8, true);
                        return;
                }
                if (strncmp(var, "memcd=", 6) == 0) {
                        do_download_ramdisk(var + 6, false);
                        return;
                }
        }
        return;
}