/*
 * 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 DMSTBI_ASSERT(x) // dummy


// @TODO: should statically initialize these for optimal thread safety
static Uint8 dm_zdefault_length[288], dm_zdefault_distance[32];


void dmZLibInit()
{
    int i;

    // Use <= to match clearly with DEFLATE spec

    // 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;
}


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)
{
    DMSTBI_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, next_code[16], sizes[17];

    memset(sizes, 0, sizeof(sizes));
    memset(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++)
    {
        next_code[i] = code;
        ctx->firstCode[i] = (Uint16) code;
        ctx->firstSymbol[i] = (Uint16) k;
        code = (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 = next_code[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 <= STBI__ZFAST_BITS)
            {
                int k = dmBitReverseN(next_code[s], s);
                while (k < STBI__ZFAST_SIZE)
                {
                    ctx->fast[k] = fastv;
                    k += (1 << s);
                }
            }
            next_code[s]++;
        }
    }

    return DMERR_OK;
}


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

    return *ctx->zbuffer++;
}


static void dmZFillBits(DMZLibContext * ctx)
{
    do
    {
        DMSTBI_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 = STBI__ZFAST_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];
    DMSTBI_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 & STBI__ZFAST_MASK];
    if (b)
    {
        int s = b >> 9;
        ctx->codeBuffer >>= s;
        ctx->numBits -= s;
        *val = b & 511;
        return DMERR_OK;
    }

    return dmZLibHuffmanDecodeSlow(ctx, huff, val);
}


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

    ctx->zout = zout;

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

    cur   = ctx->zout - ctx->zoutStart;
    limit = ctx->zoutEnd - ctx->zoutStart;

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

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

    ctx->zoutStart = newBuf;
    ctx->zout      = newBuf + cur;
    ctx->zoutEnd   = 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 *zout = ctx->zout;
    for (;;)
    {
        int z, ret;
        if ((ret = dmZLibHuffmanDecode(ctx, &ctx->zlength, &z)) != DMERR_OK)
            return ret;

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

                zout = ctx->zout;
            }
            *zout++ = (Uint8) z;
        }
        else
        {
            Uint8 *p;
            int len, dist;
            if (z == 256)
            {
                ctx->zout = zout;
                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 (zout - ctx->zoutStart < dist)
            {
                return dmErrorDBG(DMERR_DATA_ERROR,
                    "Bad Huffman block distance.\n");
            }

            if (zout + len > ctx->zoutEnd)
            {
                if ((ret = dmZLibExpand(ctx, zout, len)) != DMERR_OK)
                    return ret;
                zout = ctx->zout;
            }
            p = (Uint8 *) (zout - dist);
            if (dist == 1)
            {                   // run of one byte; common in images.
                Uint8 v = *p;
                do { *zout++ = v; } while (--len);
            }
            else
            {
                do { *zout++ = *p++; } while (--len);
            }
        }
    }
}


static const Uint8 dm_zlib_length_dezigzag[19] =
{
    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 z_codelength;
    Uint8 codeLengths[286 + 32 + 137];     // padding for maximum single op
    Uint8 codeLengthSizes[19];
    int i, n, ret;

    int hlit  = dmZReceive(ctx, 5) + 257;
    int hdist = dmZReceive(ctx, 5) + 1;
    int hclen = dmZReceive(ctx, 4) + 4;

    memset(codeLengthSizes, 0, sizeof(codeLengthSizes));

    for (i = 0; i < hclen; i++)
    {
        int s = dmZReceive(ctx, 3);
        codeLengthSizes[dm_zlib_length_dezigzag[i]] = (Uint8) s;
    }

    if ((ret = dmZLibBuildHuffmanTables(&z_codelength, codeLengthSizes, 19)) != DMERR_OK)
        return ret;

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

        DMSTBI_ASSERT(c >= 0 && c < 19);

        if (c < 16)
            codeLengths[n++] = (Uint8) c;
        else
        if (c == 16)
        {
            c = dmZReceive(ctx, 2) + 3;
            memset(codeLengths + n, codeLengths[n - 1], c);
            n += c;
        }
        else
        if (c == 17)
        {
            c = dmZReceive(ctx, 3) + 3;
            memset(codeLengths + n, 0, c);
            n += c;
        }
        else
        {
            c = dmZReceive(ctx, 7) + 11;
            memset(codeLengths + n, 0, c);
            n += c;
        }
    }

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

    if ((ret = dmZLibBuildHuffmanTables(&ctx->zlength, codeLengths, hlit)) != DMERR_OK)
        return ret;

    if ((ret = dmZLibBuildHuffmanTables(&ctx->zdistance, codeLengths + hlit, hdist)) != DMERR_OK)
        return ret;

    return DMERR_OK;
}


static int dmZLibParseUncompressedBlock(DMZLibContext * ctx)
{
    Uint8 header[4];
    int 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);

    // drain 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;
    }
    DMSTBI_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 (ctx->zbuffer + len > ctx->zbufferEnd)
    {
        return dmErrorDBG(DMERR_BOUNDS,
            "Read past buffer, probably corrupt compressed data.\n");
    }

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

    memcpy(ctx->zout, ctx->zbuffer, len);
    ctx->zbuffer += len;
    ctx->zout    += len;

    return DMERR_OK;
}


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

    ctx->numBits = 0;
    ctx->codeBuffer = 0;
    do
    {
        final = dmZReceive(ctx, 1);
        type  = dmZReceive(ctx, 2);
        if (type == 0)
        {
            if ((ret = dmZLibParseUncompressedBlock(ctx)) != DMERR_OK)
                return ret;
        }
        else
        if (type == 3)
            return DMERR_INVALID_DATA;
        else
        {
            if (type == 1)
            {
                // use fixed code lengths
                if ((ret = dmZLibBuildHuffmanTables(&ctx->zlength, dm_zdefault_length, 288)) != DMERR_OK)
                    return ret;

                if ((ret = dmZLibBuildHuffmanTables(&ctx->zdistance, dm_zdefault_distance, 32)) != 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;
}