// Copyright (c) 2016 YamaArashi

#include <stdlib.h>
#include <stdbool.h>
#include "global.h"
#include "rl.h"

unsigned char* RLDecompress(unsigned char* src, int srcSize, int* uncompressedSize) {
    if (srcSize < 4)
        goto fail;

    int destSize = (src[3] << 16) | (src[2] << 8) | src[1];

    unsigned char* dest = malloc(destSize);

    if (dest == NULL)
        goto fail;

    int srcPos = 4;
    int destPos = 0;

    for (;;) {
        if (srcPos >= srcSize)
            goto fail;

        unsigned char flags = src[srcPos++];
        bool compressed = ((flags & 0x80) != 0);

        if (compressed) {
            int length = (flags & 0x7F) + 3;
            unsigned char data = src[srcPos++];

            if (destPos + length > destSize)
                goto fail;

            for (int i = 0; i < length; i++)
                dest[destPos++] = data;
        } else {
            int length = (flags & 0x7F) + 1;

            if (destPos + length > destSize)
                goto fail;

            for (int i = 0; i < length; i++)
                dest[destPos++] = src[srcPos++];
        }

        if (destPos == destSize) {
            *uncompressedSize = destSize;
            return dest;
        }
    }

fail:
    FATAL_ERROR("Fatal error while decompressing RL file.\n");
}

unsigned char* RLCompress(unsigned char* src, int srcSize, int* compressedSize) {
    if (srcSize <= 0)
        goto fail;

    int worstCaseDestSize = 4 + srcSize * 2;

    // Round up to the next multiple of four.
    worstCaseDestSize = (worstCaseDestSize + 3) & ~3;

    unsigned char* dest = malloc(worstCaseDestSize);

    if (dest == NULL)
        goto fail;

    // header
    dest[0] = 0x30; // RL compression type
    dest[1] = (unsigned char)srcSize;
    dest[2] = (unsigned char)(srcSize >> 8);
    dest[3] = (unsigned char)(srcSize >> 16);

    int srcPos = 0;
    int destPos = 4;

    for (;;) {
        bool compress = false;
        int uncompressedStart = srcPos;
        int uncompressedLength = 0;

        while (srcPos < srcSize && uncompressedLength < (0x7F + 1)) {
            compress = (srcPos + 2 < srcSize && src[srcPos] == src[srcPos + 1] && src[srcPos] == src[srcPos + 2]);

            if (compress)
                break;

            srcPos++;
            uncompressedLength++;
        }

        if (uncompressedLength > 0) {
            dest[destPos++] = uncompressedLength - 1;

            for (int i = 0; i < uncompressedLength; i++)
                dest[destPos++] = src[uncompressedStart + i];
        }

        if (compress) {
            unsigned char data = src[srcPos];
            int compressedLength = 0;

            while (compressedLength < (0x7F + 3) && srcPos + compressedLength < srcSize &&
                   src[srcPos + compressedLength] == data) {
                compressedLength++;
            }

            dest[destPos++] = 0x80 | (compressedLength - 3);
            dest[destPos++] = data;

            srcPos += compressedLength;
        }

        if (srcPos == srcSize) {
            // Pad to multiple of 4 bytes.
            int remainder = destPos % 4;

            if (remainder != 0) {
                for (int i = 0; i < 4 - remainder; i++)
                    dest[destPos++] = 0;
            }

            *compressedSize = destPos;
            return dest;
        }
    }

fail:
    FATAL_ERROR("Fatal error while compressing RL file.\n");
}
