#include "core/compressors/ApLib.h"

#include <algorithm>

namespace bmp2tile
{

namespace
{

// ---- format constants -----------------------------------------------
// Ported from aplibsharp's constant.cs. Independently cross-checked
// against apultra's format.h, which defines the same offset thresholds
// (128/1280/32000) under different names -- two unrelated
// implementations agreeing on these exact magic numbers is strong
// confirmation both target the same real, standard aPLib bitstream.
constexpr int ThresholdByteMatchOffset = 16;      // 4 bits (0-15)
constexpr int ThresholdShortMatchOffset = 128;    // 7 bits (0-127)
constexpr int ThresholdLength2MatchOffset = 1280;
constexpr int ThresholdLength3MatchOffset = 32000;
constexpr int ThresholdOffset = 65024;  // (256-2) << 8
constexpr int ThresholdLength = 65535;

// ---- encoder state ----------------------------------------------------
// Simplified from aplibsharp's state.cs: that version also carries
// bitcost/steps/a parent pointer, needed for its backtracking optimal-
// path search. This port does a greedy left-to-right parse instead (see
// the file header comment in ApLib.h for why), so only the three fields
// that actually affect what encodings are legal/cheap right now are
// needed.
struct EncodeState
{
    int Index = 0;  // position in `input` we're about to encode from
    int Roffs = 0;  // "repeat offset": offset of the most recent length>=2 match
    int Lwm = 0;    // "last write was a match" (1 if the previous edge had length>1)
};

EncodeState Advance(const EncodeState& state, int length, int offset)
{
    EncodeState next;
    next.Index = state.Index + length;
    next.Roffs = (length > 1) ? offset : state.Roffs;
    next.Lwm = (length > 1) ? 1 : 0;
    return next;
}

// Ported from model.cs's valid(). NOTE: the reference also requires
// `offset != state.offset` (the offset of whatever edge led to this
// state) -- that's a search-pruning rule for aplibsharp's optimal-path
// search, not a format requirement (nothing in the actual bitstream
// depends on it, and the decoder has no such restriction), so it's
// intentionally not carried over to this from-scratch greedy search.
bool IsMatchValid(const EncodeState& state, int length, int offset)
{
    return (offset == state.Roffs && state.Lwm == 0 && state.Roffs != 0) ||
           (length > 3 && length <= ThresholdLength && offset > 0 && offset < ThresholdOffset) ||
           (length == 2 && offset > 0 && offset < ThresholdLength2MatchOffset) ||
           (length == 3 && offset > 0 && offset < ThresholdLength3MatchOffset);
}

// Ported from model.cs's l1valid(). Governs whether a length-1
// byte-match (offset 1-15) is legal to use right now -- specifically,
// it's excluded when it would exactly duplicate the repeat-offset match
// that's already active (state.Lwm==1 with the same offset), since that
// case is redundant with continuing the existing match instead.
bool IsL1Valid(const EncodeState& state, int offset)
{
    return offset == 0 || (offset > 0 && offset < ThresholdByteMatchOffset &&
                            (state.Lwm == 0 || (state.Lwm == 1 && offset != state.Roffs)));
}

// ---- bit/byte emission -------------------------------------------------
// Ported from tobinary.cs's putbit/putbits/putgbits and the final
// flush. The output stream interleaves a bit-packed "control" stream
// (tags, gamma-coded lengths/offset-high-bits) with raw, byte-aligned
// data (literal bytes, offset low bytes, short-match bytes) -- the
// control bits accumulate into a placeholder byte (reserved via
// PutRawByte-equivalent at construction) that gets backfilled once 8
// bits land in it, regardless of how many raw bytes get appended to the
// stream in between. Storing the placeholder's position as an index
// (not a pointer/iterator) into the shared output vector, exactly as
// aplibsharp's C# uses a List<byte> index, means this stays correct
// across any vector reallocation.
class BitWriter
{
public:
    explicit BitWriter(std::vector<uint8_t>& output) : Output(output)
    {
        BitBuffer = 1;
        BitBufferIndex = Output.size();
        Output.push_back(0);
    }

    void PutBit(int bit)
    {
        if (BitBuffer > 0xFF)
        {
            Output[BitBufferIndex] = static_cast<uint8_t>(BitBuffer & 0xFF);
            BitBufferIndex = Output.size();
            Output.push_back(0);
            BitBuffer = 1;
        }
        BitBuffer = (BitBuffer << 1) | bit;
    }

    // Ported from tobinary.cs's putgbits, kept structurally identical
    // (including the C-style `value >> msb--` idiom) to minimise the
    // chance of a subtle divergence from the reference's bit layout.
    void PutGammaBits(int value)
    {
        int msb = 15;
        while ((value >> msb--) == 0)
        {
        }

        while (msb >= 0)
        {
            const int bit = (value >> msb) & 1;
            PutBit(bit);
            --msb;
            PutBit(msb >= 0 ? 1 : 0);
        }
    }

    void PutRawByte(uint8_t b) { Output.push_back(b); }

    // Pads the final partial control byte with trailing zero bits and
    // backfills it -- must be called exactly once, after all PutBit/
    // PutGammaBits/PutRawByte calls are done.
    void Finish()
    {
        while (BitBuffer < 256)
        {
            BitBuffer <<= 1;
        }
        Output[BitBufferIndex] = static_cast<uint8_t>(BitBuffer & 0xFF);
    }

private:
    std::vector<uint8_t>& Output;
    int BitBuffer = 0;
    size_t BitBufferIndex = 0;
};

// Ported from tobinary.cs's encodestep(). Emits one literal/match edge
// and advances `state` past it.
void EmitEdge(BitWriter& writer, const std::vector<uint8_t>& input, EncodeState& state, int length, int offset)
{
    if (length == 1)
    {
        if (offset == 0 && input[static_cast<size_t>(state.Index)] != 0x00)
        {
            writer.PutBit(0);  // prefixcode_literal = {0}
            writer.PutRawByte(input[static_cast<size_t>(state.Index)]);
        }
        else
        {
            writer.PutBit(1);
            writer.PutBit(1);
            writer.PutBit(1);  // prefixcode_byte_match = {1,1,1}
            for (int s = 3; s >= 0; --s)
            {
                writer.PutBit((offset >> s) & 1);
            }
        }
    }
    else if (offset == state.Roffs && state.Lwm == 0)
    {
        // repeat offset / gapped match -- the cheapest encoding, no
        // offset bits at all.
        writer.PutBit(1);
        writer.PutBit(0);  // prefixcode_normal_match = {1,0}
        writer.PutGammaBits(2);
        writer.PutGammaBits(length);
    }
    else if (length <= 3 && offset > 0 && offset < ThresholdShortMatchOffset)
    {
        writer.PutBit(1);
        writer.PutBit(1);
        writer.PutBit(0);  // prefixcode_short_match = {1,1,0}
        writer.PutRawByte(static_cast<uint8_t>((offset << 1) | (length & 1)));
    }
    else
    {
        writer.PutBit(1);
        writer.PutBit(0);  // prefixcode_normal_match
        writer.PutGammaBits((offset >> 8) + (3 - state.Lwm));
        writer.PutRawByte(static_cast<uint8_t>(offset & 0xFF));
        const int ldo = (offset < ThresholdShortMatchOffset || offset >= ThresholdLength3MatchOffset)
                             ? 2
                             : (offset >= ThresholdLength2MatchOffset ? 1 : 0);
        writer.PutGammaBits(length - ldo);
    }

    state = Advance(state, length, offset);
}

// ---- match finding (NOT ported -- see ApLib.h) -------------------------
// A plain hash-chain finder over 2-byte prefixes, the same style already
// used for ZX7 in this project. Correctness of the *output* doesn't
// depend on this being any particular algorithm -- IsMatchValid() is
// still checked before any match this finds is used -- only the
// resulting compression ratio does.
class HashChainMatchFinder
{
public:
    explicit HashChainMatchFinder(const std::vector<uint8_t>& input)
        : Input(input), Head(1 << 16, -1), Chain(input.size(), -1)
    {
    }

    void Insert(int pos)
    {
        if (pos + 1 >= static_cast<int>(Input.size())) return;
        const int h = Hash(pos);
        Chain[static_cast<size_t>(pos)] = Head[static_cast<size_t>(h)];
        Head[static_cast<size_t>(h)] = pos;
    }

    struct Match
    {
        int Length = 0;
        int Offset = 0;
    };

    // Longest match for the bytes starting at `pos`, with offset in
    // [1, maxOffset] and length capped at maxLength. Length 0 means
    // nothing found.
    Match FindLongest(int pos, int maxOffset, int maxLength) const
    {
        Match best;
        if (pos + 1 >= static_cast<int>(Input.size()) || maxOffset < 1) return best;

        const int h = Hash(pos);
        int candidate = Head[static_cast<size_t>(h)];
        int stepsLeft = kMaxChainSteps;
        const int limit = std::min(maxLength, static_cast<int>(Input.size()) - pos);

        while (candidate >= 0 && stepsLeft-- > 0)
        {
            const int offset = pos - candidate;
            if (offset > maxOffset) break;  // chain is walked most-recent-first, so offsets only grow from here

            int length = 0;
            while (length < limit && Input[static_cast<size_t>(candidate + length)] == Input[static_cast<size_t>(pos + length)])
            {
                ++length;
            }
            if (length > best.Length)
            {
                best.Length = length;
                best.Offset = offset;
            }
            candidate = Chain[static_cast<size_t>(candidate)];
        }

        return best;
    }

private:
    int Hash(int pos) const { return Input[static_cast<size_t>(pos)] | (Input[static_cast<size_t>(pos) + 1] << 8); }

    static constexpr int kMaxChainSteps = 64;  // bounds search cost; doesn't affect correctness, only ratio/speed

    const std::vector<uint8_t>& Input;
    std::vector<int> Head;   // 65536-entry table: 2-byte hash -> most recent position with that hash, or -1
    std::vector<int> Chain;  // per-position: previous position with the same hash, or -1
};

}  // namespace

std::vector<uint8_t> CompressApLib(const std::vector<uint8_t>& input)
{
    std::vector<uint8_t> output;
    if (input.empty())
    {
        return output;
    }

    // The very first byte of an aPLib stream is always raw -- no prefix
    // bit, no framing at all (matches decode.cs's unconditional
    // `output.Add(getbyte())` before its main loop even starts).
    output.push_back(input[0]);

    BitWriter writer(output);
    EncodeState state;
    state.Index = 1;

    HashChainMatchFinder finder(input);
    finder.Insert(0);

    const int inputSize = static_cast<int>(input.size());

    while (state.Index < inputSize)
    {
        const int startIndex = state.Index;
        const int maxOffset = std::min(startIndex, ThresholdOffset - 1);
        const int maxLength = std::min(inputSize - startIndex, ThresholdLength);

        // Candidate A: continue the active repeat offset, if there is
        // one -- the cheapest possible encoding, so it's always worth
        // taking when it's at least as long as anything else found.
        int repLength = 0;
        if (state.Lwm == 0 && state.Roffs != 0 && state.Roffs <= startIndex)
        {
            while (repLength < maxLength &&
                   input[static_cast<size_t>(startIndex - state.Roffs + repLength)] == input[static_cast<size_t>(startIndex + repLength)])
            {
                ++repLength;
            }
        }

        // Candidate B: best hash-chain match at this position.
        const HashChainMatchFinder::Match found = finder.FindLongest(startIndex, maxOffset, maxLength);

        // Candidate C: length-1 byte-match (cheaper than a literal
        // whenever available -- 7 bits vs 9) or, failing that, a plain
        // literal.
        int l1Offset = 0;
        bool haveL1Match = false;
        if (input[static_cast<size_t>(startIndex)] == 0x00)
        {
            haveL1Match = true;  // the "emit 0x00 via byte-match offset 0" trick -- always legal, always cheaper
        }
        else
        {
            const int limit = std::min(startIndex, ThresholdByteMatchOffset - 1);
            for (int o = 1; o <= limit; ++o)
            {
                if (input[static_cast<size_t>(startIndex - o)] == input[static_cast<size_t>(startIndex)] && IsL1Valid(state, o))
                {
                    l1Offset = o;
                    haveL1Match = true;
                    break;  // any valid offset 1-15 costs the same, so the first one found is fine
                }
            }
        }

        int chosenLength = 1;
        int chosenOffset = haveL1Match ? l1Offset : 0;

        if (repLength >= 2 && repLength >= found.Length)
        {
            chosenLength = repLength;
            chosenOffset = state.Roffs;
        }
        else if (found.Length >= 2 && IsMatchValid(state, found.Length, found.Offset))
        {
            chosenLength = found.Length;
            chosenOffset = found.Offset;
        }
        else if (repLength >= 2)
        {
            // The hash-chain match wasn't usable (e.g. a length-2/3
            // match whose offset exceeds that length's tighter cap --
            // see IsMatchValid) but the repeat offset still works.
            chosenLength = repLength;
            chosenOffset = state.Roffs;
        }

        while (chosenLength > ThresholdLength)
        {
            EmitEdge(writer, input, state, ThresholdLength, chosenOffset);
            chosenLength -= ThresholdLength;
        }
        EmitEdge(writer, input, state, chosenLength, chosenOffset);

        for (int i = startIndex; i < state.Index; ++i)
        {
            finder.Insert(i);
        }
    }

    // End-of-stream marker: prefixcode_short_match followed by a raw
    // 0x00 byte -- decode.cs recognises this as "offset 0", which for a
    // short match is otherwise meaningless, so it's used as the
    // terminator.
    writer.PutBit(1);
    writer.PutBit(1);
    writer.PutBit(0);
    writer.PutRawByte(0x00);

    writer.Finish();

    return output;
}

}  // namespace bmp2tile
