#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

typedef unsigned char data_type;

struct Uncompressor {
    unsigned char* compressedData;
    int dataPosition;
    int varPosition;
    int bitForVarPosition;
};

data_type getByte(struct Uncompressor* uncompressor)
{
    return uncompressor->compressedData[uncompressor->dataPosition++];
}

data_type getBit(struct Uncompressor* uncompressor)
{
    if (uncompressor->bitForVarPosition == 7)
    {
        uncompressor->varPosition = uncompressor->dataPosition;
        uncompressor->dataPosition = uncompressor->varPosition + 1;
    }

    data_type bit = (uncompressor->compressedData[uncompressor->varPosition] >> uncompressor->bitForVarPosition) & 1;
    uncompressor->bitForVarPosition--;

    if (uncompressor->bitForVarPosition == -1)
    {
        uncompressor->bitForVarPosition = 7;
    }
    return bit;
}

int getInterlacedEliasGamma(struct Uncompressor* uncompressor)
{
    int value = 1;
    while (getBit(uncompressor))
    {
        value = (value << 1) | getBit(uncompressor);
    }
    return value;
}

unsigned char * realloc_output(unsigned char * output, size_t output_position, size_t * output_max_size, size_t add_data_size)
{
    if (output_position + add_data_size > *output_max_size)
    {
        *output_max_size *= 2;
        output = (unsigned char*) realloc(output, *output_max_size);
    }
    return output;
}

void uncompress(struct Uncompressor* uncompressor, size_t in_data_length, bool dsk2rom_tweak)
{
    unsigned char* output = (unsigned char*) malloc(in_data_length);// * 10 is an arbitrary number...
    size_t output_max_size = in_data_length;
    size_t output_position = 0;

    // The dsk2rom version has a fixed value of 2 for q_value and doesn't read it from the compressed variable stream
    int q_value = 2;
    if (!dsk2rom_tweak)
    {
        q_value = (getBit(uncompressor) << 2 | getBit(uncompressor) << 1 | getBit(uncompressor)) + 1;
    }

    data_type first_byte = getByte(uncompressor);
    output[output_position++] = first_byte;

    uncompressor->dataPosition = dsk2rom_tweak ? 1 : 2;

    while (uncompressor->dataPosition < in_data_length)
    {
        if (getBit(uncompressor))
        {
            int length = getInterlacedEliasGamma(uncompressor) + 1;

            if (length == 262143 || (dsk2rom_tweak && length == 131072))
            {
                break;
            }

            int offset = getByte(uncompressor);
            if (offset & 0x80)
            {
                offset = offset & 0x7F;
                switch (q_value)
                {
                        case 7:
                        offset = offset | (getBit(uncompressor) << 13);
                    case 6:
                        offset = offset | (getBit(uncompressor) << 12);
                    case 5:
                        offset = offset | (getBit(uncompressor) << 11);
                    case 4:
                        offset = offset | (getBit(uncompressor) << 10);
                    case 3:
                        offset = offset | (getBit(uncompressor) << 9);
                    case 2:
                        offset = offset | (getBit(uncompressor) << 8);
                        offset = offset | (getBit(uncompressor) << 7);
                    case 1:
                        break;
                    default:
                        printf("Invalid mode: %d\n", q_value);
                        break;
                }
                offset += 128;
            }
            offset += 1;

            output = realloc_output(output, output_position, &output_max_size, length);
            for (int i = 0; i < length; i++)
            {
                // Realloc if necessary
                if (output_position + 1 > output_max_size)
                {
                    output_max_size *= 2;
                    output = (unsigned char*) realloc(output, output_max_size);
                }

                output[output_position] = output[output_position - offset];
                output_position += 1;
            }
        }
        else
        {
            output = realloc_output(output, output_position, &output_max_size, 1);
            output[output_position++] = getByte(uncompressor);
        }
    }

    int count = 0;
    for (int i = 0; i < output_position; i++)
    {
        printf("%02x ", output[i]);
        if (++count % 30 == 0)
        {
            printf("\n");
        }
    }
    printf("\n");
    printf("Output size: %zu\n", output_position);
    free(output);
}

int main(int argc, char* argv[])
{
    if (argc < 2)
    {
        printf("Usage: %s <compressed file>\n", argv[0]);
        return 1;
    }

    FILE* file = fopen(argv[1], "rb");
    if (!file)
    {
        printf("Error reading file!\n");
        return 1;
    }

    fseek(file, 0, SEEK_END);
    long length = ftell(file);
    fseek(file, 0, SEEK_SET);

    unsigned char* compressedData = (unsigned char*) malloc(length);
    if (!fread(compressedData, length, 1, file))
    {
        printf("Error reading file!\n");
        return 1;
    }

    struct Uncompressor uncompressor = {compressedData, 0, 0, 7};
    uncompress(&uncompressor, length, true);

    free(compressedData);
    fclose(file);

    return 0;
}