/**
 * EA's RefPack compression codec.
 *
 * Red Alert 3 compresses individual `BIG4` payloads and every SAGE `.map` with
 * RefPack (the `10 FB` stream shared across EA titles). This partition decodes
 * and encodes that stream. The decoder is a port of the reference used across
 * the project (`ra3tools/ra3_big.py`, `OpenRA3`'s `ra3.fs`); the encoder is a
 * greedy LZ77 matcher that emits only canonical tokens, so anything it produces
 * is readable by the same decoder (and by the game).
 */
export module ra3.assets:refpack;

import std;

import :bytes;
import :error;

export namespace ra3::assets {

    /** Max back-reference distance (the 17-bit RefPack window). */
    inline constexpr std::size_t refpack_window = 1U << 17U;

    /** True when `data` starts with a RefPack header (`0b??010000`, `0xFB`). */
    [[nodiscard]] inline auto is_refpack(std::span<const std::byte> data) -> bool {
        return data.size() >= 2U && (std::to_integer<std::uint8_t>(data[0]) & 0x3EU) == 0x10U && std::to_integer<std::uint8_t>(data[1]) == 0xFBU;
    }

    /** Read the declared output size from a RefPack header without decompressing. */
    [[nodiscard]] inline auto refpack_output_size(std::span<const std::byte> data) -> std::uint32_t {
        if (!is_refpack(data)) throw refpack_error("not a RefPack stream");
        std::size_t pos = 0;
        const auto header = std::to_integer<std::uint8_t>(data[pos++]);
        const bool large_files = (header & 0x80U) != 0U;
        const bool compressed_size_present = (header & 0x01U) != 0U;
        pos++; // 0xFB
        const std::size_t size_bytes = large_files ? 4U : 3U;
        const auto read_size = [&]() -> std::uint32_t {
            std::uint32_t value = 0;
            for (std::size_t i = 0; i < size_bytes; ++i) {
                if (pos >= data.size()) throw refpack_error("truncated RefPack size field");
                value = (value << 8U) | std::to_integer<std::uint8_t>(data[pos++]);
            }
            return value;
        };
        if (compressed_size_present) (void) read_size();
        return read_size();
    }

    /**
     * Decompress an EA RefPack stream.
     *
     * @throws refpack_error if the stream is malformed or the length disagrees.
     */
    [[nodiscard]] inline auto refpack_decompress(std::span<const std::byte> data) -> std::vector<std::byte> {
        if (!is_refpack(data)) throw refpack_error("not a RefPack stream");

        std::size_t pos = 0;
        const auto header = std::to_integer<std::uint8_t>(data[pos++]);
        const bool large_files = (header & 0x80U) != 0U;
        const bool compressed_size_present = (header & 0x01U) != 0U;
        pos++; // 0xFB

        const std::size_t size_bytes = large_files ? 4U : 3U;
        const auto read_size = [&]() -> std::uint32_t {
            std::uint32_t value = 0;
            for (std::size_t i = 0; i < size_bytes; ++i) {
                if (pos >= data.size()) throw refpack_error("truncated RefPack size field");
                value = (value << 8U) | std::to_integer<std::uint8_t>(data[pos++]);
            }
            return value;
        };

        if (compressed_size_present) (void) read_size();
        const auto out_len = read_size();

        std::vector<std::byte> out;
        out.reserve(out_len); // guarantees no reallocation, so overlapping reads stay valid

        const auto copy_literals = [&](std::size_t count) {
            if (pos + count > data.size()) throw refpack_error("truncated RefPack literals");
            out.insert(out.end(), data.begin() + static_cast<std::ptrdiff_t>(pos), data.begin() + static_cast<std::ptrdiff_t>(pos + count));
            pos += count;
        };

        // A back-reference may overlap its own output (an RLE run): out[start + i]
        // is read one byte at a time, so the pattern repeats correctly.
        const auto copy_reference = [&](std::size_t length, std::size_t distance) {
            if (distance == 0U || distance > out.size()) throw refpack_error("RefPack back-reference out of range");
            const auto start = out.size() - distance;
            if (out.size() + length > out_len) throw refpack_error("RefPack output overrun");
            for (std::size_t i = 0; i < length; ++i) out.push_back(out[start + i]);
        };

        while (pos < data.size()) {
            const auto cmd = std::to_integer<std::uint8_t>(data[pos++]);
            if ((cmd & 0x80U) == 0U) { // 2-byte command
                if (pos >= data.size()) throw refpack_error("truncated 2-byte command");
                const auto b2 = std::to_integer<std::uint8_t>(data[pos++]);
                copy_literals(cmd & 0x03U);
                copy_reference(((cmd & 0x1CU) >> 2U) + 3U, ((cmd & 0x60U) << 3U) + b2 + 1U);
            } else if ((cmd & 0x40U) == 0U) { // 3-byte command
                if (pos + 1U >= data.size()) throw refpack_error("truncated 3-byte command");
                const auto b2 = std::to_integer<std::uint8_t>(data[pos]);
                const auto b3 = std::to_integer<std::uint8_t>(data[pos + 1U]);
                pos += 2U;
                copy_literals((b2 & 0xC0U) >> 6U);
                copy_reference((cmd & 0x3FU) + 4U, ((b2 & 0x3FU) << 8U) + b3 + 1U);
            } else if ((cmd & 0x20U) == 0U) { // 4-byte command
                if (pos + 2U >= data.size()) throw refpack_error("truncated 4-byte command");
                const auto b2 = std::to_integer<std::uint8_t>(data[pos]);
                const auto b3 = std::to_integer<std::uint8_t>(data[pos + 1U]);
                const auto b4 = std::to_integer<std::uint8_t>(data[pos + 2U]);
                pos += 3U;
                copy_literals(cmd & 0x03U);
                copy_reference(((cmd & 0x0CU) << 6U) + b4 + 5U, ((cmd & 0x10U) << 12U) + (static_cast<std::size_t>(b2) << 8U) + b3 + 1U);
            } else if (cmd < 0xFCU) { // long literal run
                copy_literals((static_cast<std::size_t>(cmd & 0x1FU) + 1U) << 2U);
            } else { // stop
                copy_literals(cmd & 0x03U);
                break;
            }
        }

        if (out.size() != out_len) throw refpack_error("RefPack length mismatch");
        return out;
    }

    /** Decompress `data` when it is RefPack, otherwise copy it unchanged. */
    [[nodiscard]] inline auto maybe_decompress(std::span<const std::byte> data) -> std::vector<std::byte> {
        if (is_refpack(data)) return refpack_decompress(data);
        return {data.begin(), data.end()};
    }

    namespace detail {
        /** True when `(length, distance)` maps to one of the three canonical tokens. */
        [[nodiscard]] inline constexpr auto refpack_match_encodable(std::size_t length, std::size_t distance) -> bool {
            if (length >= 3U && length <= 10U && distance >= 1U && distance <= 1024U) return true;
            if (length >= 4U && length <= 67U && distance >= 1U && distance <= 16384U) return true;
            if (length >= 5U && length <= 1028U && distance >= 1U && distance <= 131072U) return true;
            return false;
        }

        /** Greedy RefPack encoder state. */
        class refpack_encoder {
        public:
            explicit refpack_encoder(std::span<const std::byte> input)
                : input_(input), chain_(input.size(), -1) {
            }

            [[nodiscard]] auto run() -> std::vector<std::byte> {
                const auto n = this->input_.size();
                const bool large = n >= (1U << 24U);
                this->out_.write_u8(static_cast<std::uint8_t>(0x10U | (large ? 0x80U : 0x00U)));
                this->out_.write_u8(0xFBU);
                const std::size_t size_bytes = large ? 4U : 3U;
                for (std::size_t i = size_bytes; i-- > 0U;) this->out_.write_u8(static_cast<std::uint8_t>((n >> (8U * i)) & 0xFFU));

                std::size_t pos = 0;
                std::size_t literals_start = 0;
                while (pos < n) {
                    std::size_t best_len = 0;
                    std::size_t best_dist = 0;
                    if (pos + 2U < n) this->find_match(pos, best_len, best_dist);
                    if (best_len >= 3U) {
                        const auto pending = pos - literals_start;
                        const auto carry = pending % 4U;      // 0..3 literals ride with the token
                        const auto run = pending - carry;     // always a multiple of 4
                        this->emit_literal_run(literals_start, run);
                        this->emit_match(literals_start + run, carry, best_len, best_dist);
                        for (std::size_t i = pos; i < pos + best_len; ++i) this->insert(i);
                        pos += best_len;
                        literals_start = pos;
                    } else {
                        this->insert(pos);
                        ++pos;
                    }
                }

                const auto pending = n - literals_start;
                const auto carry = pending % 4U;
                this->emit_literal_run(literals_start, pending - carry);
                this->out_.write_u8(static_cast<std::uint8_t>(0xFCU | carry));
                this->write_literals(literals_start + (pending - carry), carry);
                return this->out_.take();
            }

        private:
            void insert(std::size_t pos) {
                if (pos + 2U >= this->input_.size()) return;
                const auto bucket = this->hash3(pos);
                this->chain_[pos] = this->head_[bucket];
                this->head_[bucket] = static_cast<std::int32_t>(pos);
            }

            void find_match(std::size_t pos, std::size_t &best_len, std::size_t &best_dist) const {
                const auto n = this->input_.size();
                const auto max_len = std::min<std::size_t>(1028U, n - pos);
                auto candidate = this->head_[this->hash3(pos)];
                int depth = 0;
                while (candidate >= 0 && depth < 64) {
                    const auto c = static_cast<std::size_t>(candidate);
                    const auto distance = pos - c;
                    if (distance > refpack_window) break; // the chain only walks backwards
                    std::size_t length = 0;
                    while (length < max_len && this->input_[c + length] == this->input_[pos + length]) ++length;
                    if (length > best_len && refpack_match_encodable(length, distance)) {
                        best_len = length;
                        best_dist = distance;
                        if (length == max_len) break;
                    }
                    candidate = this->chain_[c];
                    ++depth;
                }
            }

            void emit_literal_run(std::size_t offset, std::size_t count) {
                std::size_t remaining = count;
                std::size_t at = offset;
                while (remaining >= 4U) {
                    std::size_t step = std::min<std::size_t>(112U, remaining);
                    step -= step % 4U;
                    if (step < 4U) step = 4U;
                    this->out_.write_u8(static_cast<std::uint8_t>(0xE0U | ((step / 4U) - 1U)));
                    this->write_literals(at, step);
                    at += step;
                    remaining -= step;
                }
            }

            void emit_match(std::size_t literal_offset, std::size_t literal_count, std::size_t length, std::size_t distance) {
                const auto d = static_cast<std::uint32_t>(distance - 1U);
                if (length >= 3U && length <= 10U && distance <= 1024U) { // 2-byte token
                    const auto cmd = static_cast<std::uint8_t>((((d >> 8U) & 0x03U) << 5U) | (static_cast<std::uint32_t>(length - 3U) << 2U) |
                                                               static_cast<std::uint32_t>(literal_count));
                    this->out_.write_u8(cmd);
                    this->out_.write_u8(static_cast<std::uint8_t>(d & 0xFFU));
                } else if (length >= 4U && length <= 67U && distance <= 16384U) { // 3-byte token
                    this->out_.write_u8(static_cast<std::uint8_t>(0x80U | (length - 4U)));
                    this->out_.write_u8(static_cast<std::uint8_t>((literal_count << 6U) | ((d >> 8U) & 0x3FU)));
                    this->out_.write_u8(static_cast<std::uint8_t>(d & 0xFFU));
                } else if (length >= 5U && length <= 1028U && distance <= 131072U) { // 4-byte token
                    const auto l = static_cast<std::uint32_t>(length - 5U);
                    const auto cmd = static_cast<std::uint8_t>(0xC0U | (((l >> 8U) & 0x03U) << 2U) | ((d >> 12U) & 0x10U) |
                                                                       static_cast<std::uint32_t>(literal_count));
                    this->out_.write_u8(cmd);
                    this->out_.write_u8(static_cast<std::uint8_t>((d >> 8U) & 0xFFU));
                    this->out_.write_u8(static_cast<std::uint8_t>(d & 0xFFU));
                    this->out_.write_u8(static_cast<std::uint8_t>(l & 0xFFU));
                } else {
                    throw refpack_error("internal error: unencodable RefPack match");
                }
                this->write_literals(literal_offset, literal_count);
            }

            void write_literals(std::size_t offset, std::size_t count) {
                for (std::size_t i = 0; i < count; ++i) this->out_.write_u8(std::to_integer<std::uint8_t>(this->input_[offset + i]));
            }

            [[nodiscard]] auto hash3(std::size_t pos) const -> std::size_t {
                const auto a = static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(this->input_[pos]));
                const auto b = static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(this->input_[pos + 1U]));
                const auto c = static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(this->input_[pos + 2U]));
                return ((a | (b << 8U) | (c << 16U)) * 2654435761U) >> (32U - 16U);
            }

            std::span<const std::byte> input_;
            std::vector<std::int32_t> chain_;
            std::vector<std::int32_t> head_ = std::vector<std::int32_t>(1U << 16U, -1);
            byte_writer out_;
        };
    }

    /**
     * Compress a byte range into a canonical RefPack stream.
     *
     * The encoder is a plain greedy LZ77: it never emits a token form the
     * decoder above cannot read, so `refpack_decompress(refpack_compress(x))`
     * is lossless for every input.
     */
    [[nodiscard]] inline auto refpack_compress(std::span<const std::byte> input) -> std::vector<std::byte> {
        return detail::refpack_encoder{input}.run();
    }

} // namespace ra3::assets