// simple, minimal deflate/inflate code
// copyright 2010, Damien P. George

import "std.h";
import "iobuf.h";

uint con BSIZE = 1024;

type buf_history = (
    byte own# buf,
    uint len,
);

void buf_history:(var bh):push_byte(byte b) {
    if (bh.len + 1 <= bh.buf.alloc) {
        bh.buf[bh.len] = b;
        ++bh.len;
    } else {
        bh.buf[0 .. _] = bh.buf[1 .. bh.len];
        bh.buf[bh.len - 1] = b;
    }
}

void buf_history:(var bh):push(byte[] b) {
    if (bh.len + b.len <= bh.buf.alloc) {
        bh.buf[bh.len .. _] = b[0 .. b.len];
        bh.len += b.len;
    } else {
        uint diff = bh.len + b.len - bh.buf.alloc;
        bh.buf[0 .. _] = bh.buf[diff .. bh.len];
        bh.buf[bh.len - diff .. _] = b[0 .. b.len];
    }
}

int buf_history:(var bh):find(byte[] b) {
    for (int i = bh.len - b.len; i >= 0; --i) {
        if (str_equal(bh.buf[i, b.len], b))
            return i;
    }
    return -1;
}

void do_write_buf(obuf# obuf, byte[] b, uint back, uint len) {
    if (b.len == 1 && b[0] & 0x80 == 0) {
        obuf:write_byte(b[0]);
    } else if (b.len == 1 && back > 0 && back < 64 + 1) {
        obuf:write_byte(0x80 | (back - 1));
    } else if (b.len == 1) {
        obuf:write_byte(0xff);
        obuf:write_byte(b[0]);
    } else if (back == 0 || len == 0) {
        for (uint i = 0; i < b.len; ++i)
            do_write_buf(obuf, b[i, 1], 0, 0);
    } else if (len < 8 + 2 && back < 512 + 1) {
        --back;
        obuf:write_byte(0xe0 | (len - 2) << 1 | back >> 8);
        obuf:write_byte(back);
    } else if (len < 32 + 2) {
        obuf:write_byte(0xc0 | (len - 2));
        obuf:write_uint(back - 1);
    } else {
        obuf:write_byte(0xfe);
        obuf:write_uint(len - 2);
        obuf:write_uint(back - 1);
    }
}

void do_deflate(ibuf# ibuf, obuf# obuf) {
    var own buf = new(byte, BSIZE)[0, 0];
    buf_history bh;
    bh.buf = new(byte, BSIZE);
    bh.len = 0;

    int last_find_back = 0;
    int last_find_len = 0;
    while (!ibuf:is_eof()) {
        buf[buf.len] = ibuf:read_byte();
        ++buf.len;
        int find = bh:find(buf);
        if (find < 0) {
            if (buf.len == 1) {
                do_write_buf(obuf, buf, 0, 0);
                bh:push_byte(buf[0]);
                buf.len = 0;
            } else {
                do_write_buf(obuf, buf[0 .. buf.len - 1], last_find_back, last_find_len);
                bh:push(buf[0 .. buf.len - 1]);
                buf[0] = buf[buf.len - 1];
                buf.len = 1;
            }
            last_find_back = 0;
            last_find_len = 0;
        } else {
            last_find_back = bh.len - find;
            last_find_len = buf.len;
        }
    }
    if (buf.len > 0)
        do_write_buf(obuf, buf, last_find_back, last_find_len);
}

void do_inflate(ibuf# ibuf, obuf# obuf) {
    buf_history bh;
    bh.buf = new(byte, BSIZE * 2);
    bh.len = 0;
    while (!ibuf:is_eof()) {
        uint b = ibuf:read_byte();
        if (b & 0x80 == 0) {
            obuf:write_byte(b);
            bh:push_byte(b);
        } else if (b & 0xc0 == 0x80) {
            b = bh.buf[bh.len - (b & 0x3f) - 1];
            obuf:write_byte(b);
            bh:push_byte(b);
        } else if (b == 0xff) {
            b = ibuf:read_byte();
            obuf:write_byte(b);
            bh:push_byte(b);
        } else if (b & 0xe0 == 0xc0) {
            uint len = (b & 0x1f) + 2;
            uint back = ibuf:read_uint() + 1;
            obuf:write(bh.buf[bh.len - back, len]);
            bh:push(bh.buf[bh.len - back, len]);
        } else if (b & 0xf0 == 0xe0) {
            uint len = ((b & 0x0e) >> 1) + 2;
            uint back = (b << 8 & 0x100 | ibuf:read_uint()) + 1;
            obuf:write(bh.buf[bh.len - back, len]);
            bh:push(bh.buf[bh.len - back, len]);
        } else if (b == 0xfe) {
            uint len = ibuf:read_uint() + 2;
            uint back = ibuf:read_uint() + 1;
            obuf:write(bh.buf[bh.len - back, len]);
            bh:push(bh.buf[bh.len - back, len]);
        } else {
            assert(0);
        }
        obuf:flush();
    }
}


void usage() {
    printf("usage: comp [-/+]\n");
}

void init() {
    var own args = unix_args();
    (`do_nothing | `do_deflate | `do_inflate) do_what = `do_nothing;
    for (uint a = 1; a < args.len; ++a) {
        if (str_equal(args[a], "-")) do_what = `do_deflate;
        else if (str_equal(args[a], "+")) do_what = `do_inflate;
        else return usage();
    }
    var own ibuf = ibuf_init(0, false);
    var own obuf = obuf_init(1, false);
    cleanup { ibuf:free(); obuf:free(); }
    switch (do_what) {
        case (`do_deflate) do_deflate(ibuf, obuf);
        case (`do_inflate) do_inflate(ibuf, obuf);
    }
}