/*
 * Compressor to compress (and decompress) TGA image files using RLE compression.
 * Only works for uncompressed RGB and RLE encoded RGB.
 * Author: Wesley Stessens
 */

#ifndef COMPRESSOR_H
#define COMPRESSOR_H

#include <QString>

class Compressor
{
public:
    enum DataType {UncompressedRGB = 2, RLECompressedRGB = 10};
    static int compressOrDecompress(const QString &input, const QString &output);
    static bool compress(const QString &input, const QString &output);
    static bool decompress(const QString &input, const QString &output);
    static QString getLastError() {return lastError;}
private:
    static QString lastError;
};

#endif

/*
 * Compressor to compress (and decompress) TGA image files using RLE compression.
 * Only works for uncompressed RGB and RLE encoded RGB.
 * Author: Wesley Stessens & Toon Cuyvers
 */

#include "compressor.h"
#include <QFile>
#include <QtEndian>

QString Compressor::lastError = "Unknown";

// Returns -1 on error, 0 on successful compression, 1 on successful decompression
int Compressor::compressOrDecompress(const QString &input, const QString &output)
{
    QFile file(input);
    if (!file.open(QIODevice::ReadOnly)) {
        lastError = QString("Could not open input file ") + input +  " for reading.";
        return -1;
    }

    QByteArray rawData = file.read(16);
    int dataType = rawData[2];
    if (dataType == UncompressedRGB) {
        file.close();
        return compress(input, output) ? 0 : -1;
    } else if (dataType == RLECompressedRGB) {
        file.close();
        return decompress(input, output) ? 1 : -1;
    } else {
        file.close();
        lastError = "Input file has unsupported data type.";
        return false;
    }
}

// Author: Wesley Stessens
bool Compressor::compress(const QString &input, const QString &output)
{
    QFile file(input);
    if (!file.open(QIODevice::ReadOnly)) {
        lastError = QString("Could not open input file ") + input +  " for reading.";
        return false;
    }

    QByteArray rawData;
    int identFieldSize = file.peek(1)[0];
    rawData = file.read(identFieldSize + 18);
    int dataType = rawData[2];
    if (dataType != UncompressedRGB) {
        lastError = "Input file has unsupported data type (needs to be Uncompressed RGB).";
        file.close();
        return false;
    }
    rawData[2] = RLECompressedRGB;
    int imagePixelSize = rawData[16];
    int colorMapLength = qFromLittleEndian<qint16>(reinterpret_cast<const uchar*>(rawData.data()) + 5);
    int colorMapEntrySize = rawData[7];
    rawData.append(file.read(colorMapLength * colorMapEntrySize));

    QByteArray a = file.read(imagePixelSize >> 3);
    QByteArray b = file.read(imagePixelSize >> 3);
    while (!a.isEmpty()) {
        int num;

        if (a != b) { // Raw packets
            QByteArray colors;
            num = 0;
            while (a != b && num < 128) {
                ++num;
                colors.append(a);
                a = b;
                b = file.read(imagePixelSize >> 3);
            }
            rawData.append(num-1);
            rawData.append(colors);
        } else if (a == b) { // RLE encoded packets
            QByteArray color = a;
            num = 1;
            while (a == b && num < 128) {
                ++num;
                a = b;
                b = file.read(imagePixelSize >> 3);
            }
            rawData.append(0x80 | (num-1));
            rawData.append(color);
            a = b;
            b = file.read(imagePixelSize >> 3);
        }
    }
    file.close();

    file.setFileName(output);
    if (!file.open(QIODevice::WriteOnly)) {
        lastError = QString("Could not open output file ") + output +  " for writing.";
        return false;
    } else {
        file.write(rawData);
        file.close();
        return true;
    }
}

// Author: Wesley Stessens
bool Compressor::decompress(const QString &input, const QString &output)
{
    QFile file(input);
    if (!file.open(QIODevice::ReadOnly)) {
        lastError = QString("Could not open input file ") + input +  " for reading.";
        return false;
    }

    QByteArray rawData;
    int identFieldSize = file.peek(1)[0];
    rawData = file.read(identFieldSize + 18);
    int dataType = rawData[2];
    if (dataType != RLECompressedRGB) {
        lastError = "Input file has unsupported data type (needs to be RLE Compressed RGB).";
        file.close();
        return false;
    }
    rawData[2] = UncompressedRGB;
    int imagePixelSize = rawData[16];
    int colorMapLength = qFromLittleEndian<qint16>(reinterpret_cast<const uchar*>(rawData.data()) + 5);
    int colorMapEntrySize = rawData[7];
    rawData.append(file.read(colorMapLength * colorMapEntrySize));

    int num;
    while (!file.atEnd()) {
        QByteArray header = file.read(1);
        quint8 headerInt = static_cast<quint8>(header[0]);
        if ((headerInt & 0x80) == 0x80) { // RLE encoded packets
            num = (headerInt ^ 0x80) + 1;
            QByteArray color = file.read(imagePixelSize >> 3);
            rawData.append(color.repeated(num));
        } else { // Raw packets
            num = headerInt + 1;
            rawData.append(file.read(num * (imagePixelSize >> 3)));
        }
    }
    file.close();

    file.setFileName(output);
    if (!file.open(QIODevice::WriteOnly)) {
        lastError = QString("Could not open output file ") + output +  " for writing.";
        return false;
    } else {
        file.write(rawData);
        file.close();
        return true;
    }
}


/*
 * Main function.
 * Author: Wesley Stessens
 */

#include <QTextStream>
#include "compressor.h"

int main(int argc, char **argv)
{
    QTextStream qout(stdout);
    bool validArgs = false;
    QString action;

    if (argc == 4) {
        action = argv[1];
        action = action.trimmed().toLower();
        validArgs = action == "auto" || action == "compress" || action == "decompress";
    }

    if (!validArgs) {
        qout << "Usage: " << argv[0] << " <action> <input file> <output file>\n\taction: auto / compress / decompress\n\tinput file: TGA image to compress or decompress using RLE compression\n\toutput file: resulting TGA image after compression or decompression\n";
    } else {
        if (action == "auto") {
            int result = Compressor::compressOrDecompress(argv[2], argv[3]);
            switch (result) {
            case -1:
                qout << "could not compress/decompress file.\n";
                qout << "ERROR: " << Compressor::getLastError() << "\n";
                break;
            case 0:
                qout << "File was compressed successfully.\n";
                break;
            case 1:
                qout << "File was decompressed successfully.\n";
                break;
            }
        } else if (action == "compress") {
            bool ok = Compressor::compress(argv[2], argv[3]);
            qout << (ok ? "File was compressed successfully.\n" :
                          "Could not compress file.\n");
            if (!ok)
                qout << "ERROR: " << Compressor::getLastError() << "\n";
        } else if (action == "decompress") {
            bool ok = Compressor::decompress(argv[2], argv[3]);
            qout << (ok ? "File was decompressed successfully.\n" :
                          "Could not decompress file.\n");
            if (!ok)
                qout << "ERROR: " << Compressor::getLastError() << "\n";
        }
    }

    return 0;
}