#include "core/compressors/Stm.h"

#include <stdexcept>

namespace bmp2tile
{

namespace
{

constexpr int MinRleLen = 2;
constexpr int MaxRleLen = MinRleLen + 63;  // 65 -- exclusive loop bound, so max run length is 65
constexpr int MaxRawLen = 63;              // exclusive loop bound, so max raw run length is 63
constexpr uint8_t RleTypeNormal = 0x01;
constexpr uint8_t RleTypeIncremental = 0x03;

int HiByte(int word) { return (word >> 8) & 0xFF; }
uint8_t LoByte(int word) { return static_cast<uint8_t>(word & 0xFF); }

}  // namespace

std::vector<uint8_t> CompressStm(const std::vector<uint8_t>& input, int widthInTiles)
{
    if (widthInTiles < 0 || widthInTiles > 255)
    {
        throw std::runtime_error("STM's tilemap width header is a single byte -- width must be 0-255 tiles.");
    }

    // Reconstitute the flat 2-bytes-per-cell little-endian input into the
    // 16-bit words the reference operates on (its `buf`, a straight
    // `(unsigned short int*)source` cast that only works because the
    // reference always runs on a little-endian host -- done explicitly
    // here so it doesn't depend on host endianness).
    const size_t inSize = input.size() / 2;
    std::vector<int> buf(inSize);  // plain int, matching the promotions every read of the reference's `unsigned short int buf[]` already goes through
    for (size_t k = 0; k < inSize; ++k)
    {
        buf[k] = input[k * 2] | (input[k * 2 + 1] << 8);
    }

    std::vector<uint8_t> output;
    output.reserve(inSize + inSize / 8 + 4);  // rough guess; grows as needed regardless

    // --- state, all scoped to this one call -- the reference's
    // file-scope "globals" are actually reset at the top of every
    // compressTilemap() call, so there's no real cross-call persistence
    // to model here. ---
    size_t current = 0;
    int curHH = 0;
    int oldHH = 0;
    bool wasTempHH = false;

    auto writeRle = [&](uint8_t val, int cnt, uint8_t type)
    { output.push_back(static_cast<uint8_t>(((cnt - MinRleLen) << 2) | type));
      output.push_back(val); };

    auto writeHi = [&](int val, bool temp)
    {
        int v = val << 3;
        if (temp) v |= 0x04;
        v |= 0x02;
        output.push_back(static_cast<uint8_t>(v));
    };

    // Ported from checkHI(i): decides whether the upcoming run (of length
    // `i`, starting at `current`) needs a new high-byte control byte
    // emitted, and if so, whether that change should be treated as
    // "temporary" (silently reverted before the *next* run is
    // considered) to avoid a wasted control byte on a one-run blip back
    // to the previous high byte.
    auto checkHi = [&](int i)
    {
        if (wasTempHH)
        {
            wasTempHH = false;
            curHH = oldHH;
        }

        if (HiByte(curHH) != HiByte(buf[current]))
        {
            bool isTemp;
            if (current + static_cast<size_t>(i) == inSize)
            {
                isTemp = false;
            }
            else if (HiByte(curHH) == HiByte(buf[current + static_cast<size_t>(i)]))
            {
                isTemp = true;
                if (HiByte(buf[current + static_cast<size_t>(i) - 1]) == HiByte(buf[current + static_cast<size_t>(i)]))
                {
                    isTemp = false;
                }
            }
            else
            {
                isTemp = false;
            }

            writeHi(HiByte(buf[current]), isTemp);

            if (isTemp)
            {
                wasTempHH = true;
                oldHH = curHH;
            }
        }

        // Always remember the HH of the *last* value in this run.
        curHH = buf[current + static_cast<size_t>(i) - 1] & 0xFF00;
    };

    output.push_back(static_cast<uint8_t>(widthInTiles));  // 1-byte header: map width in tiles

    while (current < inSize)
    {
        int i = 1;

        if (current + 1 < inSize && buf[current] == buf[current + 1])
        {
            for (i = 2; i < MaxRleLen; ++i)
            {
                if (current + static_cast<size_t>(i) >= inSize) break;
                if (buf[current] != buf[current + static_cast<size_t>(i)]) break;
            }
            checkHi(i);
            writeRle(LoByte(buf[current]), i, RleTypeNormal);
        }
        else if (current + 1 < inSize && buf[current] + 1 == buf[current + 1])
        {
            for (i = 2; i < MaxRleLen; ++i)
            {
                if (current + static_cast<size_t>(i) == inSize) break;
                if (buf[current + static_cast<size_t>(i) - 1] + 1 != buf[current + static_cast<size_t>(i)]) break;
            }
            checkHi(i);
            writeRle(LoByte(buf[current]), i, RleTypeIncremental);
        }
        else
        {
            for (i = 1; i < MaxRawLen; ++i)
            {
                if (current + static_cast<size_t>(i) == inSize) break;
                if (buf[current + static_cast<size_t>(i) - 1] == buf[current + static_cast<size_t>(i)])
                {
                    --i;
                    break;
                }
                if (buf[current + static_cast<size_t>(i) - 1] + 1 == buf[current + static_cast<size_t>(i)])
                {
                    --i;
                    break;
                }
                if (HiByte(buf[current + static_cast<size_t>(i) - 1]) != HiByte(buf[current + static_cast<size_t>(i)])) break;
            }
            checkHi(i);

            output.push_back(static_cast<uint8_t>(i << 2));
            for (int j = 0; j < i; ++j)
            {
                output.push_back(LoByte(buf[current + static_cast<size_t>(j)]));
            }
        }

        current += static_cast<size_t>(i);
    }

    output.push_back(0x00);  // end-of-data marker

    return output;
}

}  // namespace bmp2tile
