#include "arch.h"

#include <assert.h>
#include <stdlib.h>
#include <string.h>
//#include <stdio.h>


enum {
    NO_CODE   = 0xffff,
    DICT_SIZE = 0xffff };


int arch_init(Arch *a) {
    memset(a, 0, sizeof(*a));
    a->entries = (Entry*)calloc(sizeof(Entry), DICT_SIZE);
    if (!a->entries) return RES_FAIL;
    for(int i = 0; i < 256; ++i) {
        Entry *e = a->entries + i;
        memset(e->children, 0xff, sizeof(e->children));
        e->value = i;
        a->root.children[i] = i;
    }
    a->nextcode = 0x100;
    a->codebits = 9;
    return RES_DONE;
}


void arch_deinit(Arch *a) {
    free(a->entries);
    memset(a, 0, sizeof(*a));
}


static void arch_entry_use(Arch *a, Entry *e) {
    if (!e->parent) return;
    if (e->next) {
        if (e->prev) e->prev->next = e->next;
                     else a->first = e->next;
        e->next->prev = e->prev;
        e->prev = a->last;
        e->next = NULL;
        a->last = a->last->next = e;
    }
    arch_entry_use(a, e->parent);
}


static Entry* arch_entry_add(Arch *a, Entry *parent, Byte value) {
    Entry *e;
    Code nextcode = a->nextcode;
    if (nextcode < DICT_SIZE) {
        ++a->nextcode;
        if (!(nextcode & a->nextcode)) ++a->codebits;
        assert(a->codebits <= 16);
        assert(nextcode < DICT_SIZE);
        e = a->entries + nextcode;
        memset(e->children, 0xff, sizeof(e->children));
        e->prev = a->last; e->next = NULL;
        if (a->last) a->last->next = e;
                     else a->first = e;
        a->last = e;
    } else {
        arch_entry_use(a, parent);

        e = a->first;
        assert(e);
        assert(e != parent);
        assert(e->parent);
        assert(e->parent->children[e->value] == e - a->entries);
        e->parent->children[e->value] = NO_CODE;

        //assert(!e->prev);
        //assert(e->next);
        //a->first = e->next; e->prev = a->last;
        //e->next = e->next->prev = NULL;
        //a->last = a->last->next = e;
    }

    assert(parent);
    assert(parent->children[value] == NO_CODE);
    e->value = value;
    e->parent = parent;
    parent->children[value] = e - a->entries;
    //printf("add code: %d\n", (int)(e - a->entries));
    //for(int i = 0; i < 256; ++i) assert(e->children[i] == NO_CODE);

    arch_entry_use(a, e);
    return e;
}


int arch_deflate(Arch *a, int final) {
    if (a->rb.bits) return RES_FAIL;
    while(1) {
        if (a->wb.end - a->wb.cur < 5) return RES_RETRY;
        Entry *e = &a->root;
        Byte *bp = a->rb.cur;
        Bits codebits = a->codebits;
        while(1) {
            Byte b;
            if (!buf_read_byte(&a->rb, &b)) {
                if (final) { arch_entry_use(a, e); break; }
                a->rb.cur = bp;
                break;
            }
            if (e->children[b] == NO_CODE) { e = arch_entry_add(a, e, b); break; }
            e = a->entries + e->children[b];
        }
        if (bp == a->rb.cur) break;

        assert(e);
        assert(e != &a->root);

        Code code; Byte b = e->value;
        if (e->parent) { code = e->parent - a->entries; }
                  else { code = a->nextcode; assert(final); }
        //printf("pos: %d/%d, code %d/%d, byte %d (%c)\n", (int)(a->rb.cur - a->rb.begin), (int)(a->rb.end - a->rb.begin), (int)code, (int)a->nextcode, (int)b, (char)b);
        if ( !buf_write_bits(&a->wb, (((Word)code) << 8) | b, codebits + 8) ) return RES_FAIL;
    }

    if (final && a->rb.cur >= a->rb.end) { buf_write_pad(&a->wb); return RES_DONE; }
    return RES_RETRY;
}


static int arch_entry_write(Arch *a, Entry *e, TinySize count) {
    if (!e->parent) {
        if (a->wb.end - a->wb.cur < count) return 0;
    } else {
        if (!arch_entry_write(a, e->parent, count+1)) return 0;
    }
    //fputc(e->value, stdout);
    buf_write_byte(&a->wb, e->value);
    return 1;
}


int arch_inflate(Arch *a, int final) {
    if (a->wb.bits) return RES_FAIL;
    while(1) {
        Byte *bp = a->rb.cur;
        Bits bb = a->rb.bits;

        Word w;
        if (!buf_read_bits(&a->rb, &w, a->codebits + 8)) break;
        Code code = w >> 8;
        Byte b = w;

        //printf("pos: %d/%d + %d, code %d/%d, byte %d (%c)\n", (int)(a->rb.cur - a->rb.begin), (int)(a->rb.end - a->rb.begin), (int)a->rb.bits, (int)code, (int)a->nextcode, (int)b, (char)b);
        if (code < a->nextcode) {
            Entry *e = &a->entries[code];
            if (!arch_entry_write(a, e, 1)) {
                a->rb.cur = bp;
                a->rb.bits = bb;
                return RES_RETRY;
            }
            //fputc(b, stdout); fputc('\n', stdout);
            if (e->children[b] == NO_CODE) e = arch_entry_add(a, e, b); else arch_entry_use(a, e);
        } else
        if (code != a->nextcode) {
            return RES_FAIL;
        } else {
            //printf("pos: %d/%d + %d, code %d/%d, byte %d (%c)\n", (int)(a->rb.cur - a->rb.begin), (int)(a->rb.end - a->rb.begin), (int)a->rb.bits, (int)code, (int)a->nextcode, (int)b, (char)b);
        }

        if (!buf_write_byte(&a->wb, b)) return RES_FAIL;
    }

    if (final) { buf_read_pad(&a->rb); return RES_DONE; }
    return RES_RETRY;
}