/*
 * DMLib
 * -- ZLib implementation
 * Public domain zlib decode v0.2 by Sean Barrett 2006-11-18
 * Modified and reformatted for DMLib by Matti 'ccr' Hamalainen
 *
 * For more information, refer to following RFCs:
 * http://tools.ietf.org/html/rfc1950 - ZLIB compressed data format
 * http://tools.ietf.org/html/rfc1951 - DEFLATE format
 *
 */
#include "dmzlib.h"


#define DM_ZLIB_TMPBUF_SIZE (16 * 1024)
#define DMZLIB_ASSERT(x) //do { printf("%s: %d\n", # x, x); } while (0) // dummy


// @TODO: should statically initialize these for optimal thread safety
static Uint8 *dm_zdefault_length = NULL,
             *dm_zdefault_distance = NULL;


int dmZLibInit()
{
    int i;

    // Use <= to match clearly with DEFLATE spec
    if ((dm_zdefault_length   = dmMalloc(DM_ZLIB_HUFF_CODES)) == NULL ||
        (dm_zdefault_distance = dmMalloc(DM_ZLIB_HUFF_DIST)) == NULL)
    {
        return dmErrorDBG(DMERR_MALLOC,
            "Failed to allocate zlib decompression tables.\n");
    }

    // Literals 0  ..143:  8 bits, codes 00110000 .. 10111111
    for (i = 0; i <= 143; i++)
        dm_zdefault_length[i] = 8;

    // Literals 144..255:  9 bits, codes 110010000 .. 111111111
    for (; i <= 255; i++)
        dm_zdefault_length[i] = 9;

    // Literals 256..279:  7 bits, codes 0000000 .. 0010111
    for (; i <= 279; i++)
        dm_zdefault_length[i] = 7;

    // Literals 280..287:  8 bits, codes 11000000 .. 11000111
    for (; i <= 287; i++)
        dm_zdefault_length[i] = 8;

    // Default distances
    for (i = 0; i <= 31; i++)
        dm_zdefault_distance[i] = 5;

    return DMERR_OK;
}


void dmZLibClose()
{
    dmFree(dm_zdefault_length);
    dm_zdefault_length = NULL;

    dmFree(dm_zdefault_distance);
    dm_zdefault_distance = NULL;
}


static inline int dmBitReverse16(int n)
{
    // "Reverse" a 16bit word through bitshifts
    n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1);
    n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2);
    n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4);
    n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8);
    return n;
}


// To bit reverse N bits, reverse 16 and shift
// e.g. 11 bits, bit reverse and shift away 5
static inline int dmBitReverseN(int v, const int bits)
{
    DMZLIB_ASSERT(bits <= 16);
    return dmBitReverse16(v) >> (16 - bits);
}


static int dmZLibBuildHuffmanTables(DMZHuffmanContext * ctx, const Uint8 * sizelist, const int num)
{
    int i, k = 0;
    int code, nextCode[16], sizes[17];

    dmMemset(sizes, 0, sizeof(sizes));
    dmMemset(ctx->fast, 0, sizeof(ctx->fast));

    for (i = 0; i < num; i++)
        sizes[sizelist[i]]++;

    sizes[0] = 0;
    for (i = 1; i < 16; i++)
    {
        if (sizes[i] > (1 << i))
        {
            return dmErrorDBG(DMERR_INTERNAL,
                "Sizes assert failed while building Huffman codes.\n");
        }
    }

    code = 0;
    for (i = 1; i < 16; i++)
    {
        nextCode[i] = code;
        ctx->firstCode[i] = (Uint16) code;
        ctx->firstSymbol[i] = (Uint16) k;
        code += sizes[i];

        if (sizes[i] && code - 1 >= (1 << i))
        {
            return dmErrorDBG(DMERR_INVALID_DATA,
                "Bad Huffman code lengths.\n");
        }

        ctx->maxCode[i] = code << (16 - i);       // preshift for inner loop
        code <<= 1;
        k += sizes[i];
    }

    ctx->maxCode[16] = 0x10000;   // sentinel
    for (i = 0; i < num; i++)
    {
        int s = sizelist[i];
        if (s)
        {
            int c = nextCode[s] - ctx->firstCode[s] + ctx->firstSymbol[s];
            Uint16 fastv = (Uint16) ((s << 9) | i);
            ctx->size[c] = (Uint8) s;
            ctx->value[c] = (Uint16) i;
            if (s <= DM_ZLIB_HFAST_BITS)
            {
                int k = dmBitReverseN(nextCode[s], s);
                while (k < DM_ZLIB_HFAST_SIZE)
                {
                    ctx->fast[k] = fastv;
                    k += (1 << s);
                }
            }
            nextCode[s]++;
        }
    }

    ctx->initialized = TRUE;
    return DMERR_OK;
}


static inline Uint8 dmZGet8(DMZLibContext * ctx)
{
    if (ctx->inBuffer >= ctx->inBufferEnd)
        return 0;

    return *ctx->inBuffer++;
}


static void dmZFillBits(DMZLibContext * ctx)
{
    do
    {
        DMZLIB_ASSERT(ctx->codeBuffer < (1U << ctx->numBits));
        ctx->codeBuffer |= dmZGet8(ctx) << ctx->numBits;
        ctx->numBits += 8;
    }
    while (ctx->numBits <= 24);
}


static inline unsigned int dmZReceive(DMZLibContext * ctx, int n)
{
    unsigned int val;

    if (ctx->numBits < n)
        dmZFillBits(ctx);

    val = ctx->codeBuffer & ((1 << n) - 1);
    ctx->codeBuffer >>= n;
    ctx->numBits -= n;

    return val;
}


static int dmZLibHuffmanDecodeSlow(DMZLibContext * ctx, DMZHuffmanContext * huff, int *val)
{
    int b, s, k;
    *val = 0;

    // not resolved by fast table, so compute it the slow way
    // use jpeg approach, which requires MSbits at top
    k = dmBitReverse16(ctx->codeBuffer);
    for (s = DM_ZLIB_HFAST_BITS + 1; ; s++)
    {
        if (k < huff->maxCode[s])
            break;
    }

    if (s == 16)
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Bad Huffman code.\n");
    }

    // code size is s, so:
    b = (k >> (16 - s)) - huff->firstCode[s] + huff->firstSymbol[s];
    DMZLIB_ASSERT(huff->size[b] == s);

    ctx->codeBuffer >>= s;
    ctx->numBits -= s;
    *val = huff->value[b];

    return DMERR_OK;
}


static inline int dmZLibHuffmanDecode(DMZLibContext * ctx, DMZHuffmanContext * huff, int *val)
{
    int b;
    if (ctx->numBits < 16)
        dmZFillBits(ctx);

    b = huff->fast[ctx->codeBuffer & DM_ZLIB_HFAST_MASK];
    if (b)
    {
        int s = b >> 9;
        ctx->codeBuffer >>= s;
        ctx->numBits -= s;
        *val = b & 511;
        return DMERR_OK;
    }

    // Not resolved by fast table, so compute it the slow way
    return dmZLibHuffmanDecodeSlow(ctx, huff, val);
}


static int dmZLibExpand(DMZLibContext * ctx, Uint8 *outBuffer, size_t n)
{
    Uint8 *newBuf;
    size_t cur, limit;

    ctx->outBuffer = outBuffer;

    if (!ctx->expandable)
    {
        return dmErrorDBG(DMERR_BOUNDS,
            "Output buffer limit hit, and is not expandable.\n");
    }

    cur   = ctx->outBuffer - ctx->outBufferStart;
    limit = ctx->outBufferEnd - ctx->outBufferStart;

    while (cur + n > limit)
        limit *= 2;

    if ((newBuf = (Uint8 *) dmRealloc(ctx->outBufferStart, limit)) == NULL)
    {
        return dmErrorDBG(DMERR_MALLOC,
            "Could not reallocate buffer.\n");
    }

    ctx->outBufferStart = newBuf;
    ctx->outBuffer      = newBuf + cur;
    ctx->outBufferEnd   = newBuf + limit;

    return DMERR_OK;
}


static const int dm_zlib_length_base[31] =
{
    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, 0, 0
};

static const int dm_zlib_length_extra[31] =
{
    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, 0, 0
};

static const int dm_zlib_dist_base[32] =
{
    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, 0, 0
};

static const int dm_zlib_dist_extra[32] =
{
    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 int dmZLibParseHuffmanBlock(DMZLibContext * ctx)
{
    Uint8 *outBuffer = ctx->outBuffer;
    for (;;)
    {
        int z, ret;
        if ((ret = dmZLibHuffmanDecode(ctx, &ctx->zlength, &z)) != DMERR_OK)
            return ret;

        if (z < 256)
        {
            if (outBuffer >= ctx->outBufferEnd)
            {
                if ((ret = dmZLibExpand(ctx, outBuffer, 1)) != DMERR_OK)
                    return ret;

                outBuffer = ctx->outBuffer;
            }
            *outBuffer++ = (Uint8) z;
        }
        else
        {
            Uint8 *ptr;
            int len, dist;
            if (z == 256)
            {
                ctx->outBuffer = outBuffer;
                return DMERR_OK;
            }

            z -= 257;

            len = dm_zlib_length_base[z];
            if (dm_zlib_length_extra[z])
                len += dmZReceive(ctx, dm_zlib_length_extra[z]);

            if ((ret = dmZLibHuffmanDecode(ctx, &ctx->zdistance, &z)) != DMERR_OK)
                return ret;

            dist = dm_zlib_dist_base[z];
            if (dm_zlib_dist_extra[z])
                dist += dmZReceive(ctx, dm_zlib_dist_extra[z]);

            if (outBuffer - ctx->outBufferStart < dist)
            {
                return dmErrorDBG(DMERR_DATA_ERROR,
                    "Bad Huffman block distance.\n");
            }

            // Check if output buffer needs to expand
            if (outBuffer + len > ctx->outBufferEnd)
            {
                if ((ret = dmZLibExpand(ctx, outBuffer, len)) != DMERR_OK)
                    return ret;
                outBuffer = ctx->outBuffer;
            }

            ptr = outBuffer - dist;
            if (dist == 1)
            {                   // run of one byte; common in images.
                Uint8 v = *ptr;
                do { *outBuffer++ = v; } while (--len);
            }
            else
            {
                do { *outBuffer++ = *ptr++; } while (--len);
            }
        }
    }
}


#define DM_ZLIB_NCODELENGTH_SIZES 19

static const Uint8 dm_zlib_length_dezigzag[DM_ZLIB_NCODELENGTH_SIZES] =
{
    16, 17, 18,  0,  8,  7,  9,
    6 , 10, 5 , 11,  4, 12,  3,
    13, 2 , 14, 1 , 15
};


static int dmZLibComputeHuffmanCodes(DMZLibContext * ctx)
{
    DMZHuffmanContext codeLengthCtx;
    Uint8 codeLengthSizes[DM_ZLIB_NCODELENGTH_SIZES];
    Uint8 codeLengths[288 + 32 + 137]; // padding for maximum single op
    int i, n, ret;

    // Get table sizes
    int hlit  = dmZReceive(ctx, 5) + 257; // 288 = X
    int hdist = dmZReceive(ctx, 5) + 1;   // 32  = X
    int hclen = dmZReceive(ctx, 4) + 4;   // 19  = DM_ZLIB_NCODELENGTH_SIZES

    if (hlit > 286 || hdist > 30)
    {
        return dmErrorDBG(DMERR_INVALID_DATA,
            "Too many Huffman length or distance symbols (%d, %d)",
            hlit, hdist);
    }

    // Get lengths table (uninitialized entries should be set to 0)
    codeLengths[256] = 0;
    dmMemset(codeLengthSizes, 0, sizeof(codeLengthSizes));
    for (i = 0; i < hclen; i++)
    {
        int s = dmZReceive(ctx, 3);
        codeLengthSizes[dm_zlib_length_dezigzag[i]] = (Uint8) s;
    }

    // Inflate the code lengths table
    if ((ret = dmZLibBuildHuffmanTables(&codeLengthCtx,
        codeLengthSizes, DM_ZLIB_NCODELENGTH_SIZES)) != DMERR_OK)
    {
        return dmErrorDBG(ret,
            "Invalid code lengths set.\n");
    }

    n = 0;
    while (n < hlit + hdist)
    {
        int c;
        if ((ret = dmZLibHuffmanDecode(ctx, &codeLengthCtx, &c)) != DMERR_OK)
            return ret;

        DMZLIB_ASSERT(c >= 0 && c < DM_ZLIB_NCODELENGTH_SIZES);

        if (c < 16)
            codeLengths[n++] = (Uint8) c;
        else
        if (c == 16)
        {
            int bv = dmZReceive(ctx, 2) + 3;
            if (n == 0)
            {
                return dmErrorDBG(DMERR_INVALID_DATA,
                    "Invalid bit length repeat.\n");
            }
            dmMemset(codeLengths + n, codeLengths[n - 1], bv);
            n += bv;
        }
        else
        if (c == 17)
        {
            int bv = dmZReceive(ctx, 3) + 3;
            dmMemset(codeLengths + n, 0, bv);
            n += bv;
        }
        else
        {
            int bv = dmZReceive(ctx, 7) + 11;
            dmMemset(codeLengths + n, 0, bv);
            n += bv;
        }
    }

    if (codeLengths[256] == 0)
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Invalid code -- missing end-of-block.\n");
    }

    if (n != hlit + hdist)
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Bad huffman codelengths.\n");
    }

    // Build the code tables
    if ((ret = dmZLibBuildHuffmanTables(&ctx->zlength, codeLengths, hlit)) != DMERR_OK)
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Invalid literal/lengths set.\n");
    }

    if ((ret = dmZLibBuildHuffmanTables(&ctx->zdistance, codeLengths + hlit, hdist)) != DMERR_OK)
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Invalid literal distances set.\n");
    }

    return DMERR_OK;
}


static int dmZLibParseUncompressedBlock(DMZLibContext * ctx)
{
    Uint8 header[4];
    size_t len, nlen, k, ret;

    // "Any bits of input up to the next byte boundary are ignored."
    if (ctx->numBits & 7)
        dmZReceive(ctx, ctx->numBits & 7);

    // Get the bit-packed data into header
    k = 0;
    while (ctx->numBits > 0)
    {
        header[k++] = (Uint8) (ctx->codeBuffer & 255);   // suppress MSVC run-time check
        ctx->codeBuffer >>= 8;
        ctx->numBits -= 8;
    }
    DMZLIB_ASSERT(ctx->numBits == 0);

    // now fill header the normal way
    while (k < 4)
        header[k++] = dmZGet8(ctx);

    len  = (header[1] << 8) | header[0];
    nlen = (header[3] << 8) | header[2];

    if (nlen != (len ^ 0xffff))
    {
        return dmErrorDBG(DMERR_DATA_ERROR,
            "Compressed data corrupt %04x :: %04x [%04x].\n",
            nlen, len, len ^ 0xffff);
    }

    if (len > 0)
    {
        if (ctx->inBuffer + len > ctx->inBufferEnd)
        {
            return dmErrorDBG(DMERR_BOUNDS,
                "Read past buffer, probably corrupt compressed data.\n");
        }

        // Check if output buffer needs to expand
        if (ctx->outBuffer + len > ctx->outBufferEnd &&
            (ret = dmZLibExpand(ctx, ctx->outBuffer, len)) != DMERR_OK)
        {
            return dmErrorDBG(DMERR_DATA_ERROR,
                "Could not expand output buffer: %d, %s\n",
                ret, dmErrorStr(ret));
        }

        // Copy uncompressed data
        memcpy(ctx->outBuffer, ctx->inBuffer, len);
        ctx->inBuffer  += len;
        ctx->outBuffer += len;
    }

    return DMERR_OK;
}


int dmZLibInitInflate(DMZLibContext *ctx)
{
    if (ctx == NULL)
        return DMERR_NULLPTR;

    dmMemset(ctx, 0, sizeof(DMZLibContext));
    return DMERR_OK;
}


void dmZLibCloseInflate(DMZLibContext *ctx)
{
    if (ctx != NULL)
    {
        dmMemset(ctx, 0, sizeof(DMZLibContext));
    }
}


int dmZLibInflate(DMZLibContext * ctx)
{
    int final, type, ret;

    ctx->numBits = 0;
    ctx->codeBuffer = 0;
    do
    {
        if (ctx->inBuffer == ctx->inBufferEnd)
            break;

        final = dmZReceive(ctx, 1);
        type  = dmZReceive(ctx, 2);
        if (type == 0)
        {
            if ((ret = dmZLibParseUncompressedBlock(ctx)) != DMERR_OK)
                return ret;
        }
        else
        if (type == 3)
        {
            return dmErrorDBG(DMERR_INVALID_DATA,
                "Invalid zlib block type.\n");
        }
        else
        {
            if (type == 1 && !ctx->zlength.initialized)
            {
                // use fixed code lengths
                if ((ret = dmZLibBuildHuffmanTables(&ctx->zlength, dm_zdefault_length, DM_ZLIB_HUFF_CODES)) != DMERR_OK)
                    return ret;

                if ((ret = dmZLibBuildHuffmanTables(&ctx->zdistance, dm_zdefault_distance, DM_ZLIB_HUFF_DIST)) != DMERR_OK)
                    return ret;
            }
            else
            if ((ret = dmZLibComputeHuffmanCodes(ctx)) != DMERR_OK)
                return ret;

            if ((ret = dmZLibParseHuffmanBlock(ctx)) != DMERR_OK)
                return ret;
        }
    }
    while (!final);

    return DMERR_OK;
}


int dmZLibParseHeader(DMZLibContext * ctx, BOOL checkPNG)
{
    // See http://tools.ietf.org/html/rfc1950
    int cmf     = dmZGet8(ctx); // Compression method and flags
    int flags   = dmZGet8(ctx); // Flags
    int cmethod = (cmf & 15);
    int cinfo   = (cmf >> 4) & 15;
    ctx->window = 1 << (8 + cinfo); // Window size (not used at the moment)

    // "The FCHECK value must be such that CMF and FLG, when viewed as
    // a 16-bit unsigned integer stored in MSB order (CMF*256 + FLG),
    // is a multiple of 31."
    if ((cmf * 256 + flags) % 31 != 0)
    {
        return dmErrorDBG(DMERR_INVALID_DATA,
            "Bad zlib header.");
    }

    // We only support compression method 8
    if (cmethod != 8)
    {
        return dmErrorDBG(DMERR_INVALID_DATA,
            "Bad or unsupported zlib compression method %d.\n",
            cmethod);
    }

    if (checkPNG && (flags & 32))
    {
        // preset dictionary not allowed in png
        return dmErrorDBG(DMERR_NOT_SUPPORTED,
            "Preset dictionary not allowed in PNG.\n");
    }
    return DMERR_OK;
}