#include <u.h>
#include <base.h>

#include <zlib.h>

typedef struct buffer
{
    union {
        struct z_stream_s; 
        z_stream z;
    };

    ubyte buf[4098];
} buffer;

// -----------------------------------------------------------------------
// Reading

typedef struct flate·Reader
{
    io·Reader rdr;
    void*     impl;

    union {
        struct buffer;
        buffer b;
    };
} flate·Reader;

flate·Reader*
flate·openreader(io·Reader rdr, void* r, mem·Allocator mem, void* m)
{
    error      err;
    flate·Reader *zrdr;

    zrdr = mem.alloc(m, 1, sizeof(*zrdr));

    zrdr->zalloc    = (void *(*)(void *, unsigned int, unsigned int))mem.alloc;
    zrdr->zfree     = mem.free;
    zrdr->opaque    = m;
    zrdr->avail_in  = rdr.read(r, 1, arrlen(zrdr->buf), zrdr->buf);
    zrdr->next_in   = zrdr->buf;

    err = inflateInit(&zrdr->b.z);

    switch (err) {
    case Z_OK:
        return zrdr;

    case Z_MEM_ERROR:
        errorf("zlib: not enough memory");
        goto ERROR;

    case Z_VERSION_ERROR:
        errorf("zlib: incompatible version");
        goto ERROR;

    case Z_STREAM_ERROR:
        errorf("zlib: incorrect input parameters");
        goto ERROR;

    default:
        errorf("zlib: unrecognized error code");
    }
ERROR:
    errorf("zlib: msg: %s", zrdr->msg);
    mem.free(m, zrdr);
    return nil;
}

error
flate·closereader(flate·Reader *rdr)
{
    int err;
    flate·Reader zrdr;

    zrdr = *rdr;
    err = inflateEnd(&zrdr.b.z);
    if (err != Z_OK) {
        errorf("zlib: failed to cleanup");
        return err;
    }
    rdr->zfree(rdr->opaque, rdr);

    return 0;
}

int
flate·read(flate·Reader *rdr, int sz, int n, void *buf)
{
    int r;
    int err;
    flate·Reader zrdr;

    zrdr = *rdr;
    zrdr.next_out  = buf;
    zrdr.avail_out = n*sz;

READ:
    err = inflate(&zrdr.b.z, Z_STREAM_END);
    switch (err) {
    case Z_OK:
        return n;

    case Z_STREAM_END:
        r  = zrdr.next_out - (ubyte*)buf;
        n -= r;
        zrdr.avail_in = zrdr.rdr.read(zrdr.impl, 1, arrlen(zrdr.buf), zrdr.buf);
        if (!zrdr.avail_in) {
            return r;
        }
        zrdr.next_in  = zrdr.buf;
        goto READ;

    case Z_NEED_DICT:
        errorf("zlib: need input dictionary");
        goto ERROR;

    case Z_STREAM_ERROR:
        errorf("zlib: inconsistent stream structure");
        goto ERROR;
    }
ERROR:
    flate·closereader(rdr);
    return -1;
}

// -----------------------------------------------------------------------
// Writing

struct flate·Writer
{
    io·Writer wtr;
    void*     impl;

    union {
        struct buffer;
        buffer b;
    };
};

flate·Writer*
flate·openwriter(io·Writer wtr, void* w, mem·Allocator mem, void* m)
{
    error      err;
    flate·Writer *zwtr;

    zwtr = mem.alloc(m, 1, sizeof(*zwtr));
    zwtr->zalloc   = (void *(*)(void *, unsigned int, unsigned int))mem.alloc;
    zwtr->zfree    = mem.free;
    zwtr->opaque   = m;
    zwtr->avail_in = 0;

    err = deflateInit(&zwtr->b.z, Z_DEFAULT_COMPRESSION);

    switch (err) {
    case Z_OK:
        return zwtr;

    case Z_MEM_ERROR:
        errorf("zlib: not enough memory");
        goto ERROR;

    case Z_VERSION_ERROR:
        errorf("zlib: incompatible version");
        goto ERROR;

    case Z_STREAM_ERROR:
        errorf("zlib: incorrect compression level");
        goto ERROR;

    default:
        errorf("zlib: unrecognized error code");
    }
ERROR:
    errorf("zlib: msg: %s", zwtr->msg);
    mem.free(m, zwtr);
    return nil;
}

error
flate·closewriter(flate·Writer *wtr)
{
    int err;
    flate·Writer zwtr;

    zwtr = *wtr;
    err = deflateEnd(&zwtr.b.z);
    if (err != Z_OK) {
        errorf("zlib: failed to cleanup");
        return err;
    }
    zwtr.zfree(zwtr.opaque, wtr);

    return 0;
}

int
flate·write(flate·Writer *wtr, int sz, int n, void *buf)
{
    int r;
    int err;
    flate·Writer zwtr;

    zwtr = *wtr;
    zwtr.next_out  = buf;
DEFLATE:
    zwtr.avail_out = n*sz;
    err = deflate(&zwtr.z, Z_NO_FLUSH);

    switch (err) {
    case Z_STREAM_END:
        return n;

    case Z_OK:
        r  = (zwtr.next_out - (ubyte*)buf)/sz;
        n -= r;
        if (!n) {
            return r;
        }
        buf += n;
        goto DEFLATE;

    case Z_STREAM_ERROR:
        errorf("zlib: bad input");
        goto ERROR;

    case Z_BUF_ERROR:
        if (!zwtr.avail_in) {
            zwtr.avail_in += zwtr.wtr.write(zwtr.impl, 1, arrlen(zwtr.buf), buf);
            if (!zwtr.avail_in) {
                errorf("reader: failed read");
                goto ERROR;
            }
            goto DEFLATE;
        }
    }

    return 0;
ERROR:
    errorf("zlib: %s", zwtr.msg);
    return -1;
}