#include "core/compressors/Zx7.h"

namespace bmp2tile
{

namespace
{

// 8-bit-limited variant: the reference zx7.h has a commented-out
// MAX_LEN=65536 for the general case, but bmp2tilecompressors builds
// against this simplified-decoder value, matching the "ZX7 (8-bit
// limited)" name bmp2tile itself uses for this plugin.
constexpr int MaxOffset = 2176;  // range 1..2176
constexpr int MaxLen = 255;      // range 2..255 (simplified decoder)

// Per-position best-known encoding, built up by Optimize() and consumed
// by EmitCompressed(). Ported field-for-field from the reference's
// `Optimal` struct.
struct Optimal
{
    size_t Bits = 0;
    int Offset = 0;
    int Len = 0;
};

int EliasGammaBits(int value)
{
    int bits = 1;
    while (value > 1)
    {
        bits += 2;
        value >>= 1;
    }
    return bits;
}

int CountBits(int offset, int len)
{
    return 1 + (offset > 128 ? 12 : 8) + EliasGammaBits(len - 1);
}

// Optimal parsing: for every position in `input`, finds the cheapest
// (fewest output bits) way to reach it via literals and back-references,
// using the same match-chain-per-2-byte-hash approach as the reference.
// bmp2tile always calls this with skip=0, so the reference's `skip`
// parameter (and the "index skipped bytes" loop it drove) is dropped here
// as dead code for our use case.
std::vector<Optimal> Optimize(const std::vector<uint8_t>& input)
{
    const size_t inputSize = input.size();

    std::vector<size_t> min(MaxOffset + 1, 0);
    std::vector<size_t> max(MaxOffset + 1, 0);
    std::vector<size_t> matches(256 * 256, 0);
    std::vector<size_t> matchSlots(inputSize, 0);
    std::vector<Optimal> optimal(inputSize);

    // First byte is always a literal.
    optimal[0].Bits = 8;

    for (size_t i = 1; i < inputSize; ++i)
    {
        optimal[i].Bits = optimal[i - 1].Bits + 9;
        const int matchIndex = (input[i - 1] << 8) | input[i];
        size_t bestLen = 1;

        for (size_t* match = &matches[static_cast<size_t>(matchIndex)]; *match != 0 && bestLen < MaxLen;
             match = &matchSlots[*match])
        {
            const int offset = static_cast<int>(i - *match);
            if (offset > MaxOffset)
            {
                *match = 0;
                break;
            }

            size_t len = 0;
            for (len = 2; len <= MaxLen && i >= len; ++len)
            {
                if (len > bestLen)
                {
                    bestLen = len;
                    const size_t bits = optimal[i - len].Bits + static_cast<size_t>(CountBits(offset, static_cast<int>(len)));
                    if (optimal[i].Bits > bits)
                    {
                        optimal[i].Bits = bits;
                        optimal[i].Offset = offset;
                        optimal[i].Len = static_cast<int>(len);
                    }
                }
                else if (max[static_cast<size_t>(offset)] != 0 && i + 1 == max[static_cast<size_t>(offset)] + len)
                {
                    len = i - min[static_cast<size_t>(offset)];
                    if (len > bestLen)
                    {
                        len = bestLen;
                    }
                }
                if (i < static_cast<size_t>(offset) + len || input[i - len] != input[i - len - static_cast<size_t>(offset)])
                {
                    break;
                }
            }
            min[static_cast<size_t>(offset)] = i + 1 - len;
            max[static_cast<size_t>(offset)] = i;
        }
        matchSlots[i] = matches[static_cast<size_t>(matchIndex)];
        matches[static_cast<size_t>(matchIndex)] = i;
    }

    return optimal;
}

// Emits the bitstream: a literal byte, or (indicator bit 1) an
// Elias-gamma length followed by a byte-or-byte-plus-4-bits offset. Bits
// are packed MSB-first into bytes, allocating a new control byte each
// time the previous one fills up. Ported directly from the reference's
// write_bit/write_byte/write_elias_gamma, minus its `diff`/delta
// bookkeeping (bmp2tile's own callers discard that value too).
class BitWriter
{
public:
    explicit BitWriter(std::vector<uint8_t>& output) : Output(output) {}

    void WriteByte(int value) { Output.push_back(static_cast<uint8_t>(value)); }

    void WriteBit(int value)
    {
        if (BitMask == 0)
        {
            BitMask = 128;
            BitIndex = Output.size();
            WriteByte(0);
        }
        if (value > 0)
        {
            Output[BitIndex] = static_cast<uint8_t>(Output[BitIndex] | BitMask);
        }
        BitMask >>= 1;
    }

    void WriteEliasGamma(int value)
    {
        int i;
        for (i = 2; i <= value; i <<= 1)
        {
            WriteBit(0);
        }
        while ((i >>= 1) > 0)
        {
            WriteBit(value & i);
        }
    }

private:
    std::vector<uint8_t>& Output;
    size_t BitIndex = 0;
    int BitMask = 0;
};

std::vector<uint8_t> EmitCompressed(std::vector<Optimal>& optimal, const std::vector<uint8_t>& input)
{
    const size_t inputSize = input.size();
    size_t inputIndex = inputSize - 1;

    std::vector<uint8_t> output;
    output.reserve((optimal[inputIndex].Bits + 18 + 7) / 8);

    // "Un-reverse" the optimal sequence. optimal[] currently holds, for
    // each position, the cumulative bit-cost of the best path *ending*
    // there (built walking forward in Optimize, but each entry only knows
    // its own cost, not what came before it in the chosen path). This
    // walks backward from the end, and -- reusing the now-unneeded .Bits
    // field as a "next position" pointer -- builds a forward linked list
    // so the emission loop below can walk start-to-end. This reuse is
    // exactly what the reference does (bits field repurposed mid-function)
    // -- preserved as-is rather than "cleaned up", since introducing a
    // separate field here would risk subtly diverging from the reference
    // logic being ported.
    optimal[inputIndex].Bits = 0;
    while (inputIndex != 0)
    {
        const size_t len = optimal[inputIndex].Len > 0 ? static_cast<size_t>(optimal[inputIndex].Len) : 1;
        const size_t inputPrev = inputIndex - len;
        optimal[inputPrev].Bits = inputIndex;
        inputIndex = inputPrev;
    }

    BitWriter writer(output);

    // First byte is always a literal.
    writer.WriteByte(input[inputIndex]);

    while ((inputIndex = optimal[inputIndex].Bits) > 0)
    {
        if (optimal[inputIndex].Len == 0)
        {
            writer.WriteBit(0);
            writer.WriteByte(input[inputIndex]);
        }
        else
        {
            writer.WriteBit(1);
            writer.WriteEliasGamma(optimal[inputIndex].Len - 1);

            int offset1 = optimal[inputIndex].Offset - 1;
            if (offset1 < 128)
            {
                writer.WriteByte(offset1);
            }
            else
            {
                offset1 -= 128;
                writer.WriteByte((offset1 & 127) | 128);
                for (int mask = 1024; mask > 127; mask >>= 1)
                {
                    writer.WriteBit(offset1 & mask);
                }
            }
        }
    }

    // End marker: a sequence indicator followed by an Elias-gamma length
    // that decodes larger than MaxLen could ever be, which the
    // decompressor recognises as "stop".
    writer.WriteBit(1);
    for (int i = 0; i < 16; ++i)
    {
        writer.WriteBit(0);
    }
    writer.WriteBit(1);

    return output;
}

}  // namespace

std::vector<uint8_t> CompressZx7(const std::vector<uint8_t>& input)
{
    if (input.empty())
    {
        return {};
    }
    std::vector<Optimal> optimal = Optimize(input);
    return EmitCompressed(optimal, input);
}

}  // namespace bmp2tile
