/*
 * Copyright (c) 2026 Netflix, Inc. Written by Warner Losh
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include "stand.h"
#include "loader_efi.h"
#include <efilib.h>
#include "decompress.h"

#include <zlib.h>
#include <bzlib.h>
#include <xz.h>
#ifdef LOADER_ZFS_SUPPORT       /* ZSTD and lzma only available with ZFS */
#include <zstd.h>
#endif
#include <sys/_param.h>

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

static EFI_MEMORY_TYPE mem_type = EfiReservedMemoryType;

struct decomp_state
{
        enum compression type;
        size_t size;            /* Best guess at final size */
        size_t alloc_size;      /* Current size of `buf` */
        size_t pages;           /* Alloc_size in pages */
        uint8_t *buf_cur;       /* Current output buffer */
        uint8_t *buf_end;       /* Current end of allocated buffer */
        EFI_PHYSICAL_ADDRESS buf; /* Decompression buffer */
        union {
                z_stream zstrm;
                bz_stream bzstrm;
                struct xz_dec *xzstrm;
#ifdef LOADER_ZFS_SUPPORT
                ZSTD_DStream *zstdstrm;
#endif
        };
        EFI_STATUS (*init)(decomp_state *dctx, uint8_t *first_buf, size_t buflen,
            size_t size_hint);
        enum step_return (*step)(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset);
        void (*fini)(decomp_state *dctx, bool flush);
};

static enum compression
what_compressed(uint8_t *buf, size_t len)
{
        /*
         * Failsafe
         */
        if (len < 32)
                return (none);
        if (memcmp(buf, "\x1f\x8b", 2) == 0) {
                printf("GZIP\n");
                return (zlib);
        }
        if (memcmp(buf, "BZh", 3) == 0) {
                printf("BZIP2\n");
                return (bzip2);
        }
        if (memcmp(buf, "\x28\xb5\x2f\xfd", 4) == 0) {
                printf("zstd\n");
                return (zstd);
        }
        if (memcmp(buf, "\xfd""7zXZ\x00", 6) == 0) {
                printf("xz\n");
                return (xz);
        }
        printf("Not compressed\n");
        return (none);
}

static EFI_STATUS
alloc_buffer(decomp_state *dctx, size_t size)
{
        dctx->alloc_size = roundup2(size, EFI_PAGE_SIZE);
        dctx->pages = dctx->alloc_size / EFI_PAGE_SIZE;
        EFI_STATUS status = BS->AllocatePages(AllocateAnyPages, mem_type, dctx->pages, &dctx->buf);
        if (EFI_ERROR(status)) {
                printf("Failed to allocate memory for %llu bytes\n", ULL(dctx->alloc_size));
                return (status);
        }
        BS->SetMem((void *)(uintptr_t)dctx->buf, dctx->alloc_size, 0);
        dctx->buf_cur = (uint8_t *)(uintptr_t)dctx->buf;
        dctx->buf_end = (uint8_t *)(uintptr_t)dctx->buf + dctx->alloc_size;
        return (EFI_SUCCESS);
}

static EFI_STATUS
grow_buffer(decomp_state *dctx)
{
        /*
         * a 1.5 exp growth trades a few more copies for a little less waste.
         */
        size_t newsz = roundup2(dctx->alloc_size * 3 / 2, EFI_PAGE_SIZE);
        size_t newpages = newsz / EFI_PAGE_SIZE;
        EFI_PHYSICAL_ADDRESS newbuf;
        EFI_STATUS status = BS->AllocatePages(AllocateAnyPages, mem_type, newpages, &newbuf);
        if (EFI_ERROR(status)) {
                printf("Failed to allocate memory for %llu bytes\n", ULL(newsz));
                return (status);
        }
        memcpy((void *)(uintptr_t)newbuf, (void *)(uintptr_t)dctx->buf, dctx->alloc_size);
        BS->FreePages(dctx->buf, dctx->pages);
        dctx->buf = newbuf;
        dctx->pages = newpages;
        dctx->buf_cur = (uint8_t *)(uintptr_t)dctx->buf + dctx->alloc_size;
        dctx->buf_end = (uint8_t *)(uintptr_t)dctx->buf + newsz;
        BS->SetMem(dctx->buf_cur, newsz - dctx->alloc_size, 0);
        dctx->alloc_size = newsz;
        return (EFI_SUCCESS);
}

static void
free_buffer(decomp_state *dctx)
{
        if (dctx->buf)
                BS->FreePages(dctx->buf, dctx->pages);
        dctx->buf = 0;
}

/*
 * zlib supprot
 */
static EFI_STATUS
zlib_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
{
        z_stream *strm = &dctx->zstrm;

        /*
         * Assume 4x compression, but start at 64MB
         */
        dctx->size = max(size_hint * 4, M(64));
        EFI_STATUS status = alloc_buffer(dctx, dctx->size);
        if (EFI_ERROR(status))
                return (status);
        memset(strm, 0, sizeof(*strm));
        strm->next_in = first_buf;
        strm->avail_in = buflen;
        return (inflateInit2(strm, 15 + 16) == Z_OK ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);
}

static enum step_return
zlib_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        z_stream *strm = &dctx->zstrm;
        size_t outlen = dctx->buf_end - dctx->buf_cur;

        strm->next_in = buf;
        strm->avail_in = len;
        strm->next_out = dctx->buf_cur;
        strm->avail_out = outlen;
        
        int ret = inflate(strm, Z_NO_FLUSH);
        dctx->buf_cur += outlen - strm->avail_out;

        if (ret == Z_STREAM_END)
                return (done);
        if (ret != Z_OK)
                return (err);
        if (dctx->buf_cur < dctx->buf_end) /* Have output space */
                return (ok);

        /*
         * We're out of space, grow the buffer and try again if there's buffer
         * space. We try again recursively since we know that will usually go
         * only 1 deep.
         */
        if (EFI_ERROR(grow_buffer(dctx)))
                return (err);
        if (strm->avail_in == 0)
                return (ok);
        size_t consumed = len - strm->avail_in;
        return (zlib_step(dctx, buf + consumed, strm->avail_in, offset + consumed));
}

static void
zlib_fini(decomp_state *dctx, bool flush)
{
        inflateEnd(&dctx->zstrm);
        if (!flush)
                return;
        free_buffer(dctx);
}

/*
 * Bzip2 supprot
 */
static EFI_STATUS
bzip2_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
{
        bz_stream *strm = &dctx->bzstrm;

        /*
         * Assume 4x compression, but start at 64MB
         */
        dctx->size = max(size_hint * 4, M(64));
        EFI_STATUS status = alloc_buffer(dctx, dctx->size);
        if (EFI_ERROR(status))
                return (status);
        memset(strm, 0, sizeof(*strm));
        strm->next_in = first_buf;
        strm->avail_in = buflen;
        return (BZ2_bzDecompressInit(strm, 0, 0) == BZ_OK ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);
}

static enum step_return
bzip2_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        bz_stream *strm = &dctx->bzstrm;
        size_t outlen = dctx->buf_end - dctx->buf_cur;

        strm->next_in = buf;
        strm->avail_in = len;
        strm->next_out = dctx->buf_cur;
        strm->avail_out = outlen;
        
        int ret = BZ2_bzDecompress(strm);
        dctx->buf_cur += outlen - strm->avail_out;

        if (ret == BZ_STREAM_END)
                return (done);
        if (ret != BZ_OK)
                return (err);
        if (dctx->buf_cur < dctx->buf_end) /* Have output space */
                return (ok);

        /*
         * We're out of space, grow the buffer and try again if there's buffer
         * space. We try again recursively since we know that will usually go
         * only 1 deep.
         */
        if (EFI_ERROR(grow_buffer(dctx)))
                return (err);
        if (strm->avail_in == 0)
                return (ok);
        size_t consumed = len - strm->avail_in;
        return (bzip2_step(dctx, buf + consumed, strm->avail_in, offset + consumed));
}

static void
bzip2_fini(decomp_state *dctx, bool flush)
{
        BZ2_bzDecompressEnd(&dctx->bzstrm);
        if (!flush)
                return;
        free_buffer(dctx);
}

/*
 * XZ support
 */
static EFI_STATUS
xz_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
{
        /*
         * Assume 4x compression, but start at 64MB
         */
        dctx->size = max(size_hint * 4, M(64));
        EFI_STATUS status = alloc_buffer(dctx, dctx->size);
        if (EFI_ERROR(status))
                return (status);
        xz_crc32_init();
        xz_crc64_init();
        dctx->xzstrm = xz_dec_init(XZ_DYNALLOC, (uint32_t)-1);
        return (dctx->xzstrm != NULL ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);     
}


static enum step_return
xz_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        struct xz_dec *strm = dctx->xzstrm;
        size_t outlen = dctx->buf_end - dctx->buf_cur;
        struct xz_buf b = { .in = buf, .in_size = len, .in_pos = 0,
                .out = dctx->buf_cur, .out_size = outlen, .out_pos = 0 };
        int ret;
        
        ret = xz_dec_run(strm, &b);
        dctx->buf_cur += b.out_pos;

        if (ret == XZ_STREAM_END)
                return (done);
        if (ret != XZ_OK) {
                switch(ret) {
                case XZ_MEM_ERROR:
                        printf("xz no memory ");
                        break;
                case XZ_DATA_ERROR:
                        printf("xz file corrupted ");
                        break;
                case XZ_FORMAT_ERROR:
                        printf("xz format not found ");
                        break;
                case XZ_OPTIONS_ERROR:
                        printf("unsupported xz option ");
                        break;
                case XZ_MEMLIMIT_ERROR:
                        printf("xz dictionary too small ");
                        break;
                default:
                        printf("xz step error %d ", ret);
                        break;
                }
                printf(" len %d offset %d\n", (int)len, (int)offset);
                return (err);
        }
        if (dctx->buf_cur < dctx->buf_end) /* Have output space */
                return (ok);

        /*
         * We're out of space, grow the buffer and try again if there's buffer
         * space. We try again recursively since we know that will usually go
         * only 1 deep.
         */
        if (EFI_ERROR(grow_buffer(dctx)))
                return (err);
        if (b.in_pos == b.in_size)
                return (ok);
        size_t consumed = b.in_pos;
        return (xz_step(dctx, buf + consumed, len - consumed, offset + consumed));
}

static void
xz_fini(decomp_state *dctx, bool flush)
{
        xz_dec_end(dctx->xzstrm);
        dctx->xzstrm = NULL;
        if (!flush)
                return;
        free_buffer(dctx);
}

/*
 * ZSTD supprot
 */
#ifdef LOADER_ZFS_SUPPORT
static EFI_STATUS
zstd_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
{
        unsigned long long size = ZSTD_getFrameContentSize(first_buf, buflen);
        if (size == ZSTD_CONTENTSIZE_ERROR)
                return (EFI_VOLUME_CORRUPTED);
        if (size == ZSTD_CONTENTSIZE_UNKNOWN)
                dctx->size = max(size_hint * 4, M(64));         /* Guess 4x compression or 64M */
        else
                dctx->size = size;                              /* We know the size */
        EFI_STATUS status = alloc_buffer(dctx, dctx->size);
        if (EFI_ERROR(status))
                return (status);

        dctx->zstdstrm = ZSTD_createDStream();
        if (dctx->zstdstrm == NULL)
                return (EFI_OUT_OF_RESOURCES);
        if (ZSTD_isError(ZSTD_initDStream(dctx->zstdstrm))) {
                ZSTD_freeDStream(dctx->zstdstrm);
                dctx->zstdstrm = NULL;
                return (EFI_OUT_OF_RESOURCES);
        }
        return (EFI_SUCCESS);
}
        
static enum step_return
zstd_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        size_t outlen = dctx->buf_end - dctx->buf_cur;
        ZSTD_inBuffer inbuf = { buf, len, 0 };
        ZSTD_outBuffer outbuf = { dctx->buf_cur, outlen, 0 };
        size_t ret;

        ret = ZSTD_decompressStream(dctx->zstdstrm, &outbuf, &inbuf);
        dctx->buf_cur += outbuf.pos;

        if (ZSTD_isError(ret))
                return (err);
        if (ret == 0)
                return (done);
        if (dctx->buf_cur < dctx->buf_end) /* Have output space */
                return (ok);

        /*
         * We're out of space, grow the buffer and try again if there's buffer
         * space. We try again recursively since we know that will usually go
         * only 1 deep.
         */
        if (EFI_ERROR(grow_buffer(dctx)))
                return (err);
        if (inbuf.size == inbuf.pos)
                return (ok);
        return (zstd_step(dctx, buf + inbuf.pos, inbuf.size - inbuf.pos, offset + inbuf.pos));
}

static void
zstd_fini(decomp_state *dctx, bool flush)
{
        ZSTD_freeDStream(dctx->zstdstrm);
        if (!flush)
                return;
        free_buffer(dctx);
}
#endif

/*
 * No / Unknown decompression fallback
 */
static EFI_STATUS
null_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
{
        dctx->size = size_hint;
        return (alloc_buffer(dctx, size_hint));
}

static enum step_return
null_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        size_t end = offset + len;
        
        if (end > dctx->size) {
                printf("Too much data recieved!");
                return (err);
        }
        if ((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf != offset) {
                printf("OH NO! The offset is %llu but I expected %llu\n", ULL(offset),
                    ULL((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf));
                return (err);
        }

        memcpy(dctx->buf_cur, buf, len);
        dctx->buf_cur += len;
        return (end == dctx->size ? done : ok);
}

static void
null_fini(decomp_state *dctx, bool flush)
{
        if (!flush)
                return;
        free_buffer(dctx);
}

decomp_state *
decomp_init(uint8_t *buf, size_t buflen, size_t size_hint)
{
        decomp_state *dctx;
        
        dctx = malloc(sizeof(*dctx));
        memset(dctx, 0, sizeof(*dctx));
        dctx->type = what_compressed(buf, buflen);
        switch (dctx->type) {
        case zlib:
                dctx->init = zlib_init;
                dctx->step = zlib_step;
                dctx->fini = zlib_fini;
                break;
        case bzip2:
                dctx->init = bzip2_init;
                dctx->step = bzip2_step;
                dctx->fini = bzip2_fini;
                break;
        case xz:
                dctx->init = xz_init;
                dctx->step = xz_step;
                dctx->fini = xz_fini;
                break;
#ifdef LOADER_ZFS_SUPPORT
        case zstd:
                dctx->init = zstd_init;
                dctx->step = zstd_step;
                dctx->fini = zstd_fini;
                break;
#endif
        case none:
                dctx->init = null_init;
                dctx->step = null_step;
                dctx->fini = null_fini;
                break;
        default:
                return (NULL);
        }

        if (EFI_ERROR(dctx->init(dctx, buf, buflen, size_hint))) {
                free(dctx);
                dctx = NULL;
        }
        return (dctx);
}

enum step_return
decomp_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
{
        return (dctx->step(dctx, buf, len, offset));
}

void
decomp_fini(decomp_state *dctx, bool flush)
{
        return (dctx->fini(dctx, flush));
}

EFI_PHYSICAL_ADDRESS
decomp_buffer(decomp_state *dctx)
{
        if (dctx == NULL)
                return (0);
        return (dctx->buf);
}

size_t
decomp_buffer_length(decomp_state *dctx)
{
        if (dctx == NULL)
                return (0);
        return ((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf);
}