// Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
// See LICENSE.md for license information.
// SPDX-License-Identifier: MIT
//
// toolZlib.c - DEFLATE for the JoeyLib host tools (toolZlib.h).
//
// Inflate decodes canonical Huffman codes a bit at a time (RFC 1951's own
// construction): small, and plenty fast for asset files. Deflate is one
// fixed-Huffman block over a greedy LZ77 parse with hash chains.

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

#include "toolZlib.h"

#define MAX_BITS            15
#define MAX_LIT_CODES       288
#define MAX_DIST_CODES      30
#define MAX_CODE_LEN_CODES  19
#define END_OF_BLOCK        256
#define FIRST_LENGTH_CODE   257
#define BLOCK_STORED        0
#define BLOCK_FIXED         1
#define BLOCK_DYNAMIC       2
#define OUT_START_BYTES     4096u
#define ADLER_MOD           65521u
#define ADLER_NMAX          5552u
#define CRC_POLY            0xEDB88320u
#define GZIP_ID1            0x1Fu
#define GZIP_ID2            0x8Bu
#define GZIP_CM_DEFLATE     8u
#define GZIP_HEADER_BYTES   10u
#define GZIP_TRAILER_BYTES  8u
#define GZIP_FHCRC          0x02u
#define GZIP_FEXTRA         0x04u
#define GZIP_FNAME          0x08u
#define GZIP_FCOMMENT       0x10u
#define ZLIB_CM_DEFLATE     8u
#define ZLIB_FDICT          0x20u
#define ZLIB_CHECK          31u
#define ZLIB_HEADER_0       0x78u
#define ZLIB_HEADER_1       0x01u
#define WINDOW_BYTES        32768u
#define HASH_BITS           15u
#define HASH_SIZE           (1u << HASH_BITS)
#define MIN_MATCH           3u
#define MAX_MATCH           258u
#define MAX_CHAIN           128u
#define NO_POS              0xFFFFFFFFu

typedef struct HuffmanT {
    uint16_t counts[MAX_BITS + 1];
    uint16_t symbols[MAX_LIT_CODES];
} HuffmanT;

typedef struct InflateT {
    const uint8_t *in;
    size_t         inLen;
    size_t         inPos;
    uint32_t       bitBuf;
    uint32_t       bitCount;
    uint8_t       *out;
    size_t         outLen;
    size_t         outCap;
    bool           error;
} InflateT;

typedef struct DeflateT {
    uint8_t *out;
    size_t   outLen;
    size_t   outCap;
    uint32_t bitBuf;
    uint32_t bitCount;
    bool     error;
} DeflateT;

static const uint16_t kLengthBase[29]                   = { 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227, 258 };
static const uint8_t  kLengthExtra[29]                  = { 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0 };
static const uint16_t kDistBase[30]                     = { 1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577 };
static const uint8_t  kDistExtra[30]                    = { 0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13 };
static const uint8_t  kCodeLenOrder[MAX_CODE_LEN_CODES] = { 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 };

static void     deflateBits(DeflateT *d, uint32_t value, uint32_t count);
static void     deflateByte(DeflateT *d, uint8_t b);
static void     deflateCode(DeflateT *d, uint32_t code, uint32_t len);
static void     deflateDistance(DeflateT *d, uint32_t dist);
static void     deflateFlush(DeflateT *d);
static void     deflateLength(DeflateT *d, uint32_t len);
static void     deflateLiteral(DeflateT *d, uint32_t sym);
static uint32_t hashAt(const uint8_t *p);
static uint32_t inflateBits(InflateT *s, uint32_t count);
static bool     inflateBlockCodes(InflateT *s, const HuffmanT *lit, const HuffmanT *dist);
static bool     inflateBuild(HuffmanT *h, const uint8_t *lengths, uint16_t n);
static int32_t  inflateDecode(InflateT *s, const HuffmanT *h);
static bool     inflateDynamic(InflateT *s);
static bool     inflateFixed(InflateT *s);
static bool     inflateOut(InflateT *s, uint8_t b);
static bool     inflateRaw(InflateT *s);
static bool     inflateStored(InflateT *s);
static uint32_t readBe32(const uint8_t *p);
static uint32_t readLe32(const uint8_t *p);
static uint32_t reverseBits(uint32_t code, uint32_t len);


static void deflateBits(DeflateT *d, uint32_t value, uint32_t count) {
    d->bitBuf   |= value << d->bitCount;
    d->bitCount += count;
    while (d->bitCount >= 8u) {
        deflateByte(d, (uint8_t)d->bitBuf);
        d->bitBuf  >>= 8;
        d->bitCount -= 8u;
    }
}


static void deflateByte(DeflateT *d, uint8_t b) {
    uint8_t *grown;

    if (d->outLen == d->outCap) {
        grown = (uint8_t *)realloc(d->out, d->outCap * 2u);
        if (grown == NULL) {
            d->error = true;
            return;
        }
        d->out     = grown;
        d->outCap *= 2u;
    }
    d->out[d->outLen++] = b;
}


// A Huffman code goes out most significant bit first.
static void deflateCode(DeflateT *d, uint32_t code, uint32_t len) {
    deflateBits(d, reverseBits(code, len), len);
}


static void deflateDistance(DeflateT *d, uint32_t dist) {
    uint32_t c;

    for (c = MAX_DIST_CODES - 1u; kDistBase[c] > dist; c--) {
    }
    deflateCode(d, c, 5u);
    deflateBits(d, dist - kDistBase[c], kDistExtra[c]);
}


static void deflateFlush(DeflateT *d) {
    if (d->bitCount > 0u) {
        deflateByte(d, (uint8_t)d->bitBuf);
        d->bitBuf   = 0u;
        d->bitCount = 0u;
    }
}


static void deflateLength(DeflateT *d, uint32_t len) {
    uint32_t c;

    for (c = 28u; kLengthBase[c] > len; c--) {
    }
    deflateLiteral(d, FIRST_LENGTH_CODE + c);
    deflateBits(d, len - kLengthBase[c], kLengthExtra[c]);
}


// The fixed literal/length code (RFC 1951 3.2.6).
static void deflateLiteral(DeflateT *d, uint32_t sym) {
    if (sym < 144u) {
        deflateCode(d, 0x30u + sym, 8u);
    } else if (sym < 256u) {
        deflateCode(d, 0x190u + (sym - 144u), 9u);
    } else if (sym < 280u) {
        deflateCode(d, sym - 256u, 7u);
    } else {
        deflateCode(d, 0xC0u + (sym - 280u), 8u);
    }
}


static uint32_t hashAt(const uint8_t *p) {
    return ((uint32_t)p[0] << 10 ^ (uint32_t)p[1] << 5 ^ (uint32_t)p[2]) & (HASH_SIZE - 1u);
}


static uint32_t inflateBits(InflateT *s, uint32_t count) {
    uint32_t value;

    while (s->bitCount < count) {
        if (s->inPos >= s->inLen) {
            s->error = true;
            return 0u;
        }
        s->bitBuf   |= (uint32_t)s->in[s->inPos++] << s->bitCount;
        s->bitCount += 8u;
    }
    value         = s->bitBuf & ((1u << count) - 1u);
    s->bitBuf   >>= count;
    s->bitCount  -= count;
    return value;
}


static bool inflateBlockCodes(InflateT *s, const HuffmanT *lit, const HuffmanT *dist) {
    int32_t  sym;
    uint32_t len;
    uint32_t back;
    uint32_t i;

    for (;;) {
        sym = inflateDecode(s, lit);
        if (sym < 0) {
            return false;
        }
        if (sym < END_OF_BLOCK) {
            if (!inflateOut(s, (uint8_t)sym)) {
                return false;
            }
            continue;
        }
        if (sym == END_OF_BLOCK) {
            return true;
        }
        sym -= FIRST_LENGTH_CODE;
        if (sym >= 29) {
            return false;
        }
        len = kLengthBase[sym] + inflateBits(s, kLengthExtra[sym]);
        sym = inflateDecode(s, dist);
        if (sym < 0 || sym >= MAX_DIST_CODES) {
            return false;
        }
        back = kDistBase[sym] + inflateBits(s, kDistExtra[sym]);
        if (s->error || back > s->outLen) {
            return false;
        }
        for (i = 0; i < len; i++) {
            if (!inflateOut(s, s->out[s->outLen - back])) {
                return false;
            }
        }
    }
}


// Canonical codes from code lengths: counts of each length, symbols in code
// order. False when the lengths over-subscribe the code space.
static bool inflateBuild(HuffmanT *h, const uint8_t *lengths, uint16_t n) {
    uint16_t offsets[MAX_BITS + 1];
    int32_t  left;
    uint16_t i;

    memset(h->counts, 0, sizeof(h->counts));
    for (i = 0; i < n; i++) {
        h->counts[lengths[i]]++;
    }
    left = 1;
    for (i = 1; i <= MAX_BITS; i++) {
        left = (left << 1) - h->counts[i];
        if (left < 0) {
            return false;
        }
    }
    offsets[1] = 0;
    for (i = 1; i < MAX_BITS; i++) {
        offsets[i + 1] = (uint16_t)(offsets[i] + h->counts[i]);
    }
    for (i = 0; i < n; i++) {
        if (lengths[i] != 0u) {
            h->symbols[offsets[lengths[i]]++] = i;
        }
    }
    return true;
}


static int32_t inflateDecode(InflateT *s, const HuffmanT *h) {
    int32_t  code;
    int32_t  first;
    int32_t  index;
    int32_t  count;
    uint32_t len;

    code  = 0;
    first = 0;
    index = 0;
    for (len = 1; len <= MAX_BITS; len++) {
        code |= (int32_t)inflateBits(s, 1u);
        if (s->error) {
            return -1;
        }
        count = h->counts[len];
        if (code - count < first) {
            return h->symbols[index + (code - first)];
        }
        index += count;
        first += count;
        first <<= 1;
        code  <<= 1;
    }
    return -1;
}


static bool inflateDynamic(InflateT *s) {
    HuffmanT lit;
    HuffmanT dist;
    uint8_t  lengths[MAX_LIT_CODES + MAX_DIST_CODES];
    uint32_t nLit;
    uint32_t nDist;
    uint32_t nCode;
    uint32_t i;
    uint32_t repeat;
    int32_t  sym;
    uint8_t  value;

    nLit  = inflateBits(s, 5u) + 257u;
    nDist = inflateBits(s, 5u) + 1u;
    nCode = inflateBits(s, 4u) + 4u;
    if (s->error || nLit > MAX_LIT_CODES || nDist > MAX_DIST_CODES) {
        return false;
    }
    memset(lengths, 0, sizeof(lengths));
    for (i = 0; i < nCode; i++) {
        lengths[kCodeLenOrder[i]] = (uint8_t)inflateBits(s, 3u);
    }
    if (!inflateBuild(&lit, lengths, MAX_CODE_LEN_CODES)) {
        return false;
    }
    i = 0;
    while (i < nLit + nDist) {
        sym = inflateDecode(s, &lit);
        if (sym < 0) {
            return false;
        }
        if (sym < 16) {
            lengths[i++] = (uint8_t)sym;
            continue;
        }
        value = 0u;
        if (sym == 16) {
            if (i == 0u) {
                return false;
            }
            value  = lengths[i - 1u];
            repeat = 3u + inflateBits(s, 2u);
        } else if (sym == 17) {
            repeat = 3u + inflateBits(s, 3u);
        } else {
            repeat = 11u + inflateBits(s, 7u);
        }
        if (s->error || i + repeat > nLit + nDist) {
            return false;
        }
        while (repeat-- > 0u) {
            lengths[i++] = value;
        }
    }
    if (lengths[END_OF_BLOCK] == 0u) {
        return false;
    }
    if (!inflateBuild(&lit, lengths, (uint16_t)nLit) || !inflateBuild(&dist, lengths + nLit, (uint16_t)nDist)) {
        return false;
    }
    return inflateBlockCodes(s, &lit, &dist);
}


static bool inflateFixed(InflateT *s) {
    HuffmanT lit;
    HuffmanT dist;
    uint8_t  lengths[MAX_LIT_CODES];
    uint16_t i;

    for (i = 0; i < 144u; i++) {
        lengths[i] = 8u;
    }
    for (; i < 256u; i++) {
        lengths[i] = 9u;
    }
    for (; i < 280u; i++) {
        lengths[i] = 7u;
    }
    for (; i < MAX_LIT_CODES; i++) {
        lengths[i] = 8u;
    }
    inflateBuild(&lit, lengths, MAX_LIT_CODES);
    for (i = 0; i < MAX_DIST_CODES; i++) {
        lengths[i] = 5u;
    }
    inflateBuild(&dist, lengths, MAX_DIST_CODES);
    return inflateBlockCodes(s, &lit, &dist);
}


static bool inflateOut(InflateT *s, uint8_t b) {
    uint8_t *grown;

    if (s->outLen == s->outCap) {
        grown = (uint8_t *)realloc(s->out, s->outCap * 2u);
        if (grown == NULL) {
            return false;
        }
        s->out     = grown;
        s->outCap *= 2u;
    }
    s->out[s->outLen++] = b;
    return true;
}


static bool inflateRaw(InflateT *s) {
    uint32_t last;
    uint32_t type;
    bool     ok;

    do {
        last = inflateBits(s, 1u);
        type = inflateBits(s, 2u);
        if (s->error) {
            return false;
        }
        if (type == BLOCK_STORED) {
            ok = inflateStored(s);
        } else if (type == BLOCK_FIXED) {
            ok = inflateFixed(s);
        } else if (type == BLOCK_DYNAMIC) {
            ok = inflateDynamic(s);
        } else {
            ok = false;
        }
        if (!ok) {
            return false;
        }
    } while (last == 0u);
    // The wrapper's trailer starts at the next whole byte.
    s->bitBuf   = 0u;
    s->bitCount = 0u;
    return true;
}


static bool inflateStored(InflateT *s) {
    uint32_t len;
    uint32_t nlen;

    s->bitBuf   = 0u;
    s->bitCount = 0u;
    if (s->inPos + 4u > s->inLen) {
        return false;
    }
    len   = (uint32_t)s->in[s->inPos] | (uint32_t)s->in[s->inPos + 1u] << 8;
    nlen  = (uint32_t)s->in[s->inPos + 2u] | (uint32_t)s->in[s->inPos + 3u] << 8;
    s->inPos += 4u;
    if ((len ^ 0xFFFFu) != nlen || s->inPos + len > s->inLen) {
        return false;
    }
    while (len-- > 0u) {
        if (!inflateOut(s, s->in[s->inPos++])) {
            return false;
        }
    }
    return true;
}


static uint32_t readBe32(const uint8_t *p) {
    return (uint32_t)p[0] << 24 | (uint32_t)p[1] << 16 | (uint32_t)p[2] << 8 | (uint32_t)p[3];
}


static uint32_t readLe32(const uint8_t *p) {
    return (uint32_t)p[0] | (uint32_t)p[1] << 8 | (uint32_t)p[2] << 16 | (uint32_t)p[3] << 24;
}


static uint32_t reverseBits(uint32_t code, uint32_t len) {
    uint32_t out;
    uint32_t i;

    out = 0u;
    for (i = 0; i < len; i++) {
        out  = (out << 1) | (code & 1u);
        code >>= 1;
    }
    return out;
}


uint32_t toolAdler32(uint32_t seed, const uint8_t *data, size_t len) {
    uint32_t a;
    uint32_t b;
    size_t   chunk;

    a = seed & 0xFFFFu;
    b = seed >> 16;
    while (len > 0u) {
        chunk = len < ADLER_NMAX ? len : ADLER_NMAX;
        len  -= chunk;
        while (chunk-- > 0u) {
            a += *data++;
            b += a;
        }
        a %= ADLER_MOD;
        b %= ADLER_MOD;
    }
    return b << 16 | a;
}


uint32_t toolCrc32(uint32_t seed, const uint8_t *data, size_t len) {
    static uint32_t table[256];
    static bool     built = false;
    uint32_t        c;
    uint32_t        n;
    uint32_t        k;

    if (!built) {
        for (n = 0; n < 256u; n++) {
            c = n;
            for (k = 0; k < 8u; k++) {
                c = (c & 1u) != 0u ? CRC_POLY ^ (c >> 1) : c >> 1;
            }
            table[n] = c;
        }
        built = true;
    }
    c = seed ^ 0xFFFFFFFFu;
    while (len-- > 0u) {
        c = table[(c ^ *data++) & 0xFFu] ^ (c >> 8);
    }
    return c ^ 0xFFFFFFFFu;
}


bool toolDeflateZlib(const uint8_t *in, size_t inLen, uint8_t **out, size_t *outLen) {
    DeflateT  d;
    uint32_t *head;
    uint32_t *prev;
    uint32_t  adler;
    uint32_t  h;
    uint32_t  cand;
    uint32_t  chain;
    uint32_t  bestLen;
    uint32_t  bestDist;
    uint32_t  len;
    uint32_t  limit;
    size_t    pos;
    size_t    k;

    memset(&d, 0, sizeof(d));
    d.outCap = OUT_START_BYTES;
    d.out    = (uint8_t *)malloc(d.outCap);
    head     = (uint32_t *)malloc(HASH_SIZE * sizeof(uint32_t));
    prev     = (uint32_t *)malloc(WINDOW_BYTES * sizeof(uint32_t));
    if (d.out == NULL || head == NULL || prev == NULL) {
        free(d.out);
        free(head);
        free(prev);
        return false;
    }
    for (k = 0; k < HASH_SIZE; k++) {
        head[k] = NO_POS;
    }
    deflateByte(&d, ZLIB_HEADER_0);
    deflateByte(&d, ZLIB_HEADER_1);
    deflateBits(&d, 1u, 1u);
    deflateBits(&d, BLOCK_FIXED, 2u);
    pos = 0;
    while (pos < inLen) {
        bestLen  = 0u;
        bestDist = 0u;
        if (pos + MIN_MATCH <= inLen) {
            h     = hashAt(in + pos);
            cand  = head[h];
            limit = (uint32_t)(inLen - pos < MAX_MATCH ? inLen - pos : MAX_MATCH);
            for (chain = 0; cand != NO_POS && chain < MAX_CHAIN && pos - cand <= WINDOW_BYTES; chain++) {
                for (len = 0; len < limit && in[cand + len] == in[pos + len]; len++) {
                }
                if (len > bestLen) {
                    bestLen  = len;
                    bestDist = (uint32_t)(pos - cand);
                    if (len == limit) {
                        break;
                    }
                }
                cand = prev[cand % WINDOW_BYTES];
            }
        }
        if (bestLen >= MIN_MATCH) {
            deflateLength(&d, bestLen);
            deflateDistance(&d, bestDist);
        } else {
            bestLen = 1u;
            deflateLiteral(&d, in[pos]);
        }
        for (k = 0; k < bestLen; k++, pos++) {
            if (pos + MIN_MATCH <= inLen) {
                h                        = hashAt(in + pos);
                prev[pos % WINDOW_BYTES] = head[h];
                head[h]                  = (uint32_t)pos;
            }
        }
    }
    deflateLiteral(&d, END_OF_BLOCK);
    deflateFlush(&d);
    adler = toolAdler32(1u, in, inLen);
    deflateByte(&d, (uint8_t)(adler >> 24));
    deflateByte(&d, (uint8_t)(adler >> 16));
    deflateByte(&d, (uint8_t)(adler >> 8));
    deflateByte(&d, (uint8_t)adler);
    free(head);
    free(prev);
    if (d.error) {
        free(d.out);
        return false;
    }
    *out    = d.out;
    *outLen = d.outLen;
    return true;
}


bool toolInflate(const uint8_t *in, size_t inLen, ToolZlibFormatE fmt, uint8_t **out, size_t *outLen) {
    InflateT s;
    size_t   pos;
    uint8_t  flags;
    uint32_t xlen;

    pos = 0;
    if (fmt == TOOL_ZLIB_AUTO) {
        fmt = (inLen >= 2u && in[0] == GZIP_ID1 && in[1] == GZIP_ID2) ? TOOL_ZLIB_GZIP : TOOL_ZLIB_ZLIB;
    }
    if (fmt == TOOL_ZLIB_ZLIB) {
        if (inLen < 6u || (in[0] & 0x0Fu) != ZLIB_CM_DEFLATE || ((uint32_t)in[0] << 8 | in[1]) % ZLIB_CHECK != 0u || (in[1] & ZLIB_FDICT) != 0u) {
            return false;
        }
        pos = 2u;
    } else if (fmt == TOOL_ZLIB_GZIP) {
        if (inLen < GZIP_HEADER_BYTES + GZIP_TRAILER_BYTES || in[0] != GZIP_ID1 || in[1] != GZIP_ID2 || in[2] != GZIP_CM_DEFLATE) {
            return false;
        }
        flags = in[3];
        pos   = GZIP_HEADER_BYTES;
        if ((flags & GZIP_FEXTRA) != 0u) {
            if (pos + 2u > inLen) {
                return false;
            }
            xlen = (uint32_t)in[pos] | (uint32_t)in[pos + 1u] << 8;
            pos += 2u + xlen;
        }
        if ((flags & GZIP_FNAME) != 0u) {
            while (pos < inLen && in[pos] != 0u) {
                pos++;
            }
            pos++;
        }
        if ((flags & GZIP_FCOMMENT) != 0u) {
            while (pos < inLen && in[pos] != 0u) {
                pos++;
            }
            pos++;
        }
        if ((flags & GZIP_FHCRC) != 0u) {
            pos += 2u;
        }
        if (pos >= inLen) {
            return false;
        }
    }
    memset(&s, 0, sizeof(s));
    s.in     = in;
    s.inLen  = inLen;
    s.inPos  = pos;
    s.outCap = OUT_START_BYTES;
    s.out    = (uint8_t *)malloc(s.outCap);
    if (s.out == NULL) {
        return false;
    }
    if (!inflateRaw(&s)) {
        free(s.out);
        return false;
    }
    if (fmt == TOOL_ZLIB_ZLIB) {
        if (s.inPos + 4u > inLen || readBe32(in + s.inPos) != toolAdler32(1u, s.out, s.outLen)) {
            free(s.out);
            return false;
        }
    } else if (fmt == TOOL_ZLIB_GZIP) {
        if (s.inPos + GZIP_TRAILER_BYTES > inLen || readLe32(in + s.inPos) != toolCrc32(0u, s.out, s.outLen) || readLe32(in + s.inPos + 4u) != (uint32_t)s.outLen) {
            free(s.out);
            return false;
        }
    }
    *out    = s.out;
    *outLen = s.outLen;
    return true;
}