// SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
// SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
// Licensed under OWL v1.0+. See LICENSE.

#include <stdlib.h>
#include <string.h>

#include "miniz.h"

#include <lauxlib.h>
#include <lua.h>

#define DEFAULT_LEVEL    6
#define MAX_LEVEL        10
#define DEFAULT_MAX_SIZE (64 * 1024 * 1024) /* 64 MB decompression bomb guard */

#define RETURN_ERR(L, msg)      \
    do {                        \
        lua_pushnil(L);         \
        lua_pushstring(L, msg); \
        return 2;               \
    } while (0)

/* ── helpers ────────────────────────────────────────────────────────── */

static int get_level(lua_State *L, int idx) {
    int level = luaL_optint(L, idx, DEFAULT_LEVEL);
    if (level < 0)
        level = 0;
    if (level > MAX_LEVEL)
        level = MAX_LEVEL;
    return level;
}

static size_t get_max_size(lua_State *L, int idx) {
    lua_Number n = luaL_optnumber(L, idx, (lua_Number)DEFAULT_MAX_SIZE);
    if (n <= 0)
        return DEFAULT_MAX_SIZE;
    if (n > (lua_Number)DEFAULT_MAX_SIZE)
        return DEFAULT_MAX_SIZE;
    return (size_t)n;
}

/* ── deflate (raw) ──────────────────────────────────────────────────── */

static int lua_zlib_deflate(lua_State *L) {
    size_t src_len;
    const char *src = luaL_checklstring(L, 1, &src_len);
    int level       = get_level(L, 2);

    mz_ulong bound     = mz_deflateBound(NULL, (mz_ulong)src_len);
    unsigned char *out = (unsigned char *)malloc(bound);
    if (!out)
        RETURN_ERR(L, "memory allocation failed");

    mz_stream stream;
    memset(&stream, 0, sizeof(stream));
    int ret = mz_deflateInit2(&stream, level, MZ_DEFLATED, -MZ_DEFAULT_WINDOW_BITS, 9, MZ_DEFAULT_STRATEGY);
    if (ret != MZ_OK) {
        free(out);
        RETURN_ERR(L, "deflate init failed");
    }

    stream.next_in   = (const unsigned char *)src;
    stream.avail_in  = (unsigned int)src_len;
    stream.next_out  = out;
    stream.avail_out = (unsigned int)bound;

    ret = mz_deflate(&stream, MZ_FINISH);
    if (ret != MZ_STREAM_END) {
        mz_deflateEnd(&stream);
        free(out);
        RETURN_ERR(L, "deflate failed");
    }

    lua_pushlstring(L, (const char *)out, stream.total_out);
    mz_deflateEnd(&stream);
    free(out);
    return 1;
}

/* ── compress (zlib-wrapped) ────────────────────────────────────────── */

static int lua_zlib_compress(lua_State *L) {
    size_t src_len;
    const char *src = luaL_checklstring(L, 1, &src_len);
    int level       = get_level(L, 2);

    mz_ulong dest_len   = mz_compressBound((mz_ulong)src_len);
    unsigned char *dest = (unsigned char *)malloc(dest_len);
    if (!dest)
        RETURN_ERR(L, "memory allocation failed");

    int ret = mz_compress2(dest, &dest_len, (const unsigned char *)src, (mz_ulong)src_len, level);
    if (ret != MZ_OK) {
        free(dest);
        RETURN_ERR(L, "compress failed");
    }

    lua_pushlstring(L, (const char *)dest, dest_len);
    free(dest);
    return 1;
}

/* ── gzip ───────────────────────────────────────────────────────────── */

static int lua_zlib_gzip(lua_State *L) {
    size_t src_len;
    const char *src = luaL_checklstring(L, 1, &src_len);
    int level       = get_level(L, 2);

    /* Raw deflate first */
    mz_ulong bound          = mz_deflateBound(NULL, (mz_ulong)src_len);
    unsigned char *deflated = (unsigned char *)malloc(bound);
    if (!deflated)
        RETURN_ERR(L, "memory allocation failed");

    mz_stream stream;
    memset(&stream, 0, sizeof(stream));
    int ret = mz_deflateInit2(&stream, level, MZ_DEFLATED, -MZ_DEFAULT_WINDOW_BITS, 9, MZ_DEFAULT_STRATEGY);
    if (ret != MZ_OK) {
        free(deflated);
        RETURN_ERR(L, "deflate init failed");
    }

    stream.next_in   = (const unsigned char *)src;
    stream.avail_in  = (unsigned int)src_len;
    stream.next_out  = deflated;
    stream.avail_out = (unsigned int)bound;

    ret = mz_deflate(&stream, MZ_FINISH);
    if (ret != MZ_STREAM_END) {
        mz_deflateEnd(&stream);
        free(deflated);
        RETURN_ERR(L, "deflate failed");
    }
    mz_ulong deflated_len = stream.total_out;
    mz_deflateEnd(&stream);

    /* Build gzip: 10-byte header + deflated data + 8-byte trailer (CRC32 + ISIZE) */
    size_t gzip_len         = 10 + deflated_len + 8;
    unsigned char *gzip_buf = (unsigned char *)malloc(gzip_len);
    if (!gzip_buf) {
        free(deflated);
        RETURN_ERR(L, "memory allocation failed");
    }

    /* Gzip header (RFC 1952 minimal) */
    gzip_buf[0] = 0x1f; /* ID1 */
    gzip_buf[1] = 0x8b; /* ID2 */
    gzip_buf[2] = 0x08; /* CM = deflate */
    gzip_buf[3] = 0x00; /* FLG = none */
    gzip_buf[4] = 0x00; /* MTIME */
    gzip_buf[5] = 0x00;
    gzip_buf[6] = 0x00;
    gzip_buf[7] = 0x00;
    gzip_buf[8] = 0x00; /* XFL */
    gzip_buf[9] = 0xff; /* OS = unknown */

    memcpy(gzip_buf + 10, deflated, deflated_len);
    free(deflated);

    /* Trailer: CRC32 + original size (both little-endian) */
    mz_ulong crc           = mz_crc32(MZ_CRC32_INIT, (const unsigned char *)src, src_len);
    unsigned char *trailer = gzip_buf + 10 + deflated_len;
    trailer[0]             = (unsigned char)(crc & 0xff);
    trailer[1]             = (unsigned char)((crc >> 8) & 0xff);
    trailer[2]             = (unsigned char)((crc >> 16) & 0xff);
    trailer[3]             = (unsigned char)((crc >> 24) & 0xff);
    mz_ulong isize         = (mz_ulong)(src_len & 0xffffffff);
    trailer[4]             = (unsigned char)(isize & 0xff);
    trailer[5]             = (unsigned char)((isize >> 8) & 0xff);
    trailer[6]             = (unsigned char)((isize >> 16) & 0xff);
    trailer[7]             = (unsigned char)((isize >> 24) & 0xff);

    lua_pushlstring(L, (const char *)gzip_buf, gzip_len);
    free(gzip_buf);
    return 1;
}

/* ── inflate (raw) ──────────────────────────────────────────────────── */

static int lua_zlib_inflate(lua_State *L) {
    size_t src_len;
    const char *src = luaL_checklstring(L, 1, &src_len);
    size_t max_size = get_max_size(L, 2);

    size_t out_len = src_len * 4;
    if (out_len < 256)
        out_len = 256;
    if (out_len > max_size)
        out_len = max_size;

    unsigned char *out = (unsigned char *)malloc(out_len);
    if (!out)
        RETURN_ERR(L, "memory allocation failed");

    mz_stream stream;
    memset(&stream, 0, sizeof(stream));
    int ret = mz_inflateInit2(&stream, -MZ_DEFAULT_WINDOW_BITS);
    if (ret != MZ_OK) {
        free(out);
        RETURN_ERR(L, "inflate init failed");
    }

    stream.next_in   = (const unsigned char *)src;
    stream.avail_in  = (unsigned int)src_len;
    stream.next_out  = out;
    stream.avail_out = (unsigned int)out_len;

    while (1) {
        ret = mz_inflate(&stream, MZ_NO_FLUSH);
        if (ret == MZ_STREAM_END)
            break;
        if (ret == MZ_BUF_ERROR || (ret == MZ_OK && stream.avail_out == 0)) {
            size_t new_len = out_len * 2;
            if (new_len > max_size) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "decompressed size exceeds max_size limit");
            }
            unsigned char *new_out = (unsigned char *)realloc(out, new_len);
            if (!new_out) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "memory allocation failed");
            }
            out              = new_out;
            stream.next_out  = out + out_len;
            stream.avail_out = (unsigned int)(new_len - out_len);
            out_len          = new_len;
            continue;
        }
        if (ret != MZ_OK) {
            mz_inflateEnd(&stream);
            free(out);
            RETURN_ERR(L, "inflate failed");
        }
    }

    lua_pushlstring(L, (const char *)out, stream.total_out);
    lua_pushinteger(L, (lua_Integer)stream.total_in);
    mz_inflateEnd(&stream);
    free(out);
    return 2;
}

/* ── decompress (zlib-wrapped) ──────────────────────────────────────── */

static int lua_zlib_decompress(lua_State *L) {
    size_t src_len;
    const char *src = luaL_checklstring(L, 1, &src_len);
    size_t max_size = get_max_size(L, 2);

    size_t out_len = src_len * 4;
    if (out_len < 256)
        out_len = 256;
    if (out_len > max_size)
        out_len = max_size;

    unsigned char *out = (unsigned char *)malloc(out_len);
    if (!out)
        RETURN_ERR(L, "memory allocation failed");

    mz_stream stream;
    memset(&stream, 0, sizeof(stream));
    int ret = mz_inflateInit2(&stream, MZ_DEFAULT_WINDOW_BITS);
    if (ret != MZ_OK) {
        free(out);
        RETURN_ERR(L, "decompress init failed");
    }

    stream.next_in   = (const unsigned char *)src;
    stream.avail_in  = (unsigned int)src_len;
    stream.next_out  = out;
    stream.avail_out = (unsigned int)out_len;

    while (1) {
        ret = mz_inflate(&stream, MZ_NO_FLUSH);
        if (ret == MZ_STREAM_END)
            break;
        if (ret == MZ_BUF_ERROR || (ret == MZ_OK && stream.avail_out == 0)) {
            size_t new_len = out_len * 2;
            if (new_len > max_size) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "decompressed size exceeds max_size limit");
            }
            unsigned char *new_out = (unsigned char *)realloc(out, new_len);
            if (!new_out) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "memory allocation failed");
            }
            out              = new_out;
            stream.next_out  = out + out_len;
            stream.avail_out = (unsigned int)(new_len - out_len);
            out_len          = new_len;
            continue;
        }
        if (ret != MZ_OK) {
            mz_inflateEnd(&stream);
            free(out);
            RETURN_ERR(L, "decompress failed");
        }
    }

    lua_pushlstring(L, (const char *)out, stream.total_out);
    lua_pushinteger(L, (lua_Integer)stream.total_in);
    mz_inflateEnd(&stream);
    free(out);
    return 2;
}

/* ── gunzip ─────────────────────────────────────────────────────────── */

static int lua_zlib_gunzip(lua_State *L) {
    size_t src_len;
    const unsigned char *src = (const unsigned char *)luaL_checklstring(L, 1, &src_len);
    size_t max_size          = get_max_size(L, 2);

    /* Parse gzip header (RFC 1952) */
    if (src_len < 18)
        RETURN_ERR(L, "input too short for gzip");
    if (src[0] != 0x1f || src[1] != 0x8b)
        RETURN_ERR(L, "invalid gzip magic");
    if (src[2] != 0x08)
        RETURN_ERR(L, "unsupported gzip compression method");

    unsigned char flg = src[3];
    size_t pos        = 10;

    /* FEXTRA */
    if (flg & 0x04) {
        if (pos + 2 > src_len)
            RETURN_ERR(L, "truncated gzip header (FEXTRA)");
        size_t xlen = src[pos] | ((size_t)src[pos + 1] << 8);
        pos += 2 + xlen;
    }
    /* FNAME */
    if (flg & 0x08) {
        while (pos < src_len && src[pos] != 0)
            pos++;
        pos++; /* skip null terminator */
    }
    /* FCOMMENT */
    if (flg & 0x10) {
        while (pos < src_len && src[pos] != 0)
            pos++;
        pos++; /* skip null terminator */
    }
    /* FHCRC */
    if (flg & 0x02) {
        pos += 2;
    }

    if (pos >= src_len)
        RETURN_ERR(L, "truncated gzip header");

    /* Compressed data ends 8 bytes before the end (CRC32 + ISIZE trailer) */
    if (src_len - pos < 8)
        RETURN_ERR(L, "truncated gzip trailer");
    size_t compressed_len           = src_len - pos - 8;
    const unsigned char *compressed = src + pos;
    const unsigned char *trailer    = src + src_len - 8;

    /* Raw inflate */
    size_t out_len = compressed_len * 4;
    if (out_len < 256)
        out_len = 256;
    if (out_len > max_size)
        out_len = max_size;

    unsigned char *out = (unsigned char *)malloc(out_len);
    if (!out)
        RETURN_ERR(L, "memory allocation failed");

    mz_stream stream;
    memset(&stream, 0, sizeof(stream));
    int ret = mz_inflateInit2(&stream, -MZ_DEFAULT_WINDOW_BITS);
    if (ret != MZ_OK) {
        free(out);
        RETURN_ERR(L, "inflate init failed");
    }

    stream.next_in   = compressed;
    stream.avail_in  = (unsigned int)compressed_len;
    stream.next_out  = out;
    stream.avail_out = (unsigned int)out_len;

    while (1) {
        ret = mz_inflate(&stream, MZ_NO_FLUSH);
        if (ret == MZ_STREAM_END)
            break;
        if (ret == MZ_BUF_ERROR || (ret == MZ_OK && stream.avail_out == 0)) {
            size_t new_len = out_len * 2;
            if (new_len > max_size) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "decompressed size exceeds max_size limit");
            }
            unsigned char *new_out = (unsigned char *)realloc(out, new_len);
            if (!new_out) {
                mz_inflateEnd(&stream);
                free(out);
                RETURN_ERR(L, "memory allocation failed");
            }
            out              = new_out;
            stream.next_out  = out + out_len;
            stream.avail_out = (unsigned int)(new_len - out_len);
            out_len          = new_len;
            continue;
        }
        if (ret != MZ_OK) {
            mz_inflateEnd(&stream);
            free(out);
            RETURN_ERR(L, "gunzip inflate failed");
        }
    }

    size_t decompressed_len = stream.total_out;
    mz_inflateEnd(&stream);

    /* Validate CRC32 trailer */
    mz_ulong expected_crc = (mz_ulong)trailer[0] | ((mz_ulong)trailer[1] << 8) | ((mz_ulong)trailer[2] << 16) |
                            ((mz_ulong)trailer[3] << 24);
    mz_ulong actual_crc = mz_crc32(MZ_CRC32_INIT, out, decompressed_len);
    if (actual_crc != expected_crc) {
        free(out);
        RETURN_ERR(L, "gzip CRC32 mismatch");
    }

    /* Validate ISIZE trailer */
    mz_ulong expected_isize = (mz_ulong)trailer[4] | ((mz_ulong)trailer[5] << 8) | ((mz_ulong)trailer[6] << 16) |
                              ((mz_ulong)trailer[7] << 24);
    if ((decompressed_len & 0xffffffff) != expected_isize) {
        free(out);
        RETURN_ERR(L, "gzip ISIZE mismatch");
    }

    lua_pushlstring(L, (const char *)out, decompressed_len);
    lua_pushinteger(L, (lua_Integer)(pos + stream.total_in + 8));
    free(out);
    return 2;
}

/* ── crc32 ──────────────────────────────────────────────────────────── */

static int lua_zlib_crc32(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    mz_ulong init    = (mz_ulong)luaL_optinteger(L, 2, MZ_CRC32_INIT);
    mz_ulong result  = mz_crc32(init, (const unsigned char *)data, len);
    lua_pushnumber(L, (lua_Number)result);
    return 1;
}

/* ── adler32 ────────────────────────────────────────────────────────── */

static int lua_zlib_adler32(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    mz_ulong init    = (mz_ulong)luaL_optinteger(L, 2, 1); /* adler32 initial value is 1 */
    mz_ulong result  = mz_adler32(init, (const unsigned char *)data, len);
    lua_pushnumber(L, (lua_Number)result);
    return 1;
}

/* ── module registration ────────────────────────────────────────────── */

static luaL_Reg funcs[] = {
    {"deflate",    lua_zlib_deflate   },
    {"compress",   lua_zlib_compress  },
    {"gzip",       lua_zlib_gzip      },
    {"inflate",    lua_zlib_inflate   },
    {"decompress", lua_zlib_decompress},
    {"gunzip",     lua_zlib_gunzip    },
    {"crc32",      lua_zlib_crc32     },
    {"adler32",    lua_zlib_adler32   },
    {NULL,         NULL               }
};

int luaopen_zlib_core(lua_State *L) {
    luaL_newlib(L, funcs);
    return 1;
}