#include <fstream>
#include <iomanip>
#include <iostream>
#include <vector>

using data_type = unsigned char;

class Uncompressor
{
public:
    explicit Uncompressor(const std::vector<unsigned char>& inData, bool dsk2rom_tweak = false)
        : compressedData(inData)
    {
        std::vector<unsigned char> output;

        // 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() << 2 | getBit() << 1 | getBit()) + 1;
        }

        data_type first_byte = getByte();
        output.push_back(first_byte);
        dataPosition = dsk2rom_tweak ? 1 : 2;

        while (dataPosition < inData.size())
        {
            if (getBit())// Back reference
            {
                int length = getInterlacedEliasGamma() + 1;

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

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

                for (int i = 0; i < length; i++)
                {
                    output.push_back(output[output.size() - offset]);
                }
            }
            else// Literal data
            {
                output.push_back(getByte());
            }
        }

        int count = 0;
        for (unsigned char byte: output)
        {
            std::cout << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(byte) << " ";
            if (++count % 30 == 0)
            {
                std::cout << std::endl;
            }
        }

        std::cout << std::endl;
        std::cout << "Output size: " << std::dec << static_cast<int>(output.size()) << std::endl;
    }

private:
    int dataPosition = 0;
    int varPosition = 0;
    int bitForVarPosition = 7;

    const std::vector<unsigned char>& compressedData;

    data_type getByte()
    {
        return compressedData[dataPosition++];
    }

    data_type getBit()
    {
        if (bitForVarPosition == 7)
        {
            varPosition = dataPosition;
            dataPosition = varPosition + 1;
        }

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

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

    int getInterlacedEliasGamma()
    {
        int value = 1;
        while (getBit())
        {
            value = (value << 1) | getBit();
        }
        return value;
    }
};

int main(int argc, char* argv[])
{
    if (argc < 2)
    {
        std::cout << "Usage: " << argv[0] << " <compressed file>" << std::endl;
        return 1;
    }

    std::ifstream file(argv[1], std::ios::binary | std::ios::ate);
    if (!file)
    {
        std::cout << "Error reading file!" << std::endl;
        return 1;
    }

    std::streamsize length = file.tellg();
    file.seekg(0, std::ios::beg);

    std::vector<unsigned char> compressedData(length);
    if (!file.read(reinterpret_cast<char*>(compressedData.data()), length))
    {
        std::cout << "Error reading file!" << std::endl;
        return 1;
    }

    Uncompressor data(compressedData);

    return 0;
}