#include "bitwriter.h"
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <assert.h>

struct BitWriter
{
    FILE *file;
    uint8_t byte;
    int mod;
};

BitWriter *BitWriter_init(FILE *file)
{
    BitWriter *bw = malloc(sizeof(BitWriter));
    assert(bw);
    bw->file = file;
    bw->byte = 0;
    bw->mod = 0;
    return bw;
}

void BitWriter_free(BitWriter *bw)
{
    free(bw);
}

void BitWriter_write_bit(BitWriter *bw, bool bit)
{
    bw->byte |= (bit > 0) << (bw->mod++);

    if (bw->mod == 8)
    {
        fwrite(&bw->byte, 1, 1, bw->file);
        bw->byte = 0;
        bw->mod = 0;
    }
}

void BitWriter_flush(BitWriter *bw)
{
    while (bw->mod != 0) BitWriter_write_bit(bw, 0);
}

void BitWriter_write_bin(BitWriter *bw, uint32_t bits, int number_to_write)
{
    for (int b = number_to_write - 1; b >= 0; b--)
        BitWriter_write_bit(bw, bits & (1 << b));
}

void BitWriter_write_bin_reverse(BitWriter *bw, uint32_t bits, int number_to_write)
{
    for (int b = 0; b < number_to_write; b++)
        BitWriter_write_bit(bw, bits & (1 << b));
}

void BitWriter_write_alpha(BitWriter *bw, uint32_t byte)
{
    assert(byte <= 285);

    if (byte <= 143) /* 8 bit encoding */
    {
        uint32_t to_encode = 48 + (uint32_t) byte;
        BitWriter_write_bin(bw, to_encode, 8);
    } 
    else if (byte <= 255) /* 9 bit encoding */
    {
        uint32_t to_encode = 256 + (uint32_t) byte;
        BitWriter_write_bin(bw, to_encode, 9);
    } 
    else if (byte <= 279) /* 7 bit encoding */
    {
        uint32_t to_encode = (uint32_t) byte - 256;
        BitWriter_write_bin(bw, to_encode, 7);
    } 
    else if (byte <= 287)
    {
        uint32_t to_encode = (uint32_t) byte - 88;
        BitWriter_write_bin(bw, to_encode, 8);
    }
}

void BitWriter_write_length(BitWriter *bw, uint32_t length)
{
    assert(3 <= length && length <= 258);

    if (length <= 10) /* 3 - 10 */
    {
        BitWriter_write_alpha(bw, length + 254);
    }
    else if (length <= 18) /* 11 - 22 */
    {
        BitWriter_write_alpha(bw, (length + 519) >> 1);
        BitWriter_write_bin_reverse(bw, 1 + length, 1);
    }
    else if (length <= 34) /* 23 - 34 */
    {
        BitWriter_write_alpha(bw, (length + 1057) >> 2);
        BitWriter_write_bin_reverse(bw, 1 + length, 2);
    }
    else if (length <= 66) /* 35 - 66 */
    {
        BitWriter_write_alpha(bw, (length + 2149) >> 3);
        BitWriter_write_bin_reverse(bw, 5 + length, 3);
    }
    else if (length <= 130) /* 67 - 130 */
    {
        BitWriter_write_alpha(bw, (length + 4365) >> 4);
        BitWriter_write_bin_reverse(bw, 13 + length, 4);
    }
    else if (length <= 257) /* 131 - 257 */
    {
        BitWriter_write_alpha(bw, (length + 8861) >> 5);
        BitWriter_write_bin_reverse(bw, 29 + length, 5);
    }
    else if (length == 258)
    {
        BitWriter_write_alpha(bw, 285);
    }
    else
    {
        assert(0);
    }
}

void BitWriter_write_distance(BitWriter *bw, uint32_t dist)
{
    assert(1 <= dist && dist <= 32768);

    if (dist <= 4) BitWriter_write_bin(bw, dist - 1, 5);
    else
    {
        int bits = 1;

        while (true)
        {
            if (dist <= 1 << (bits + 2))
            {
                BitWriter_write_bin(bw, (dist + bits * (1 << (bits + 1)) - 1) >> bits, 5);
                BitWriter_write_bin_reverse(bw, dist - 1, bits);
                return;
            } 
            else bits++;
        }
    }
}