/*
 * Copyright (c) 2022, Michiel Visser <opensource@webmichiel.nl>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include <AK/BitStream.h>
#include <AK/CircularQueue.h>
#include <AK/FixedArray.h>
#include <AK/Vector.h>

namespace Compress {

namespace Brotli {

class CanonicalCode {
public:
    CanonicalCode() = default;
    CanonicalCode(Vector<size_t> codes, Vector<size_t> values)
        : m_symbol_codes(move(codes))
        , m_symbol_values(move(values)) {};

    static ErrorOr<CanonicalCode> read_prefix_code(LittleEndianInputBitStream&, size_t alphabet_size);
    static ErrorOr<CanonicalCode> read_simple_prefix_code(LittleEndianInputBitStream&, size_t alphabet_size);
    static ErrorOr<CanonicalCode> read_complex_prefix_code(LittleEndianInputBitStream&, size_t alphabet_size, size_t hskip);

    ErrorOr<size_t> read_symbol(LittleEndianInputBitStream&) const;

private:
    static ErrorOr<size_t> read_complex_prefix_code_length(LittleEndianInputBitStream&);

    Vector<size_t> m_symbol_codes;
    Vector<size_t> m_symbol_values;
};

}

class BrotliDecompressionStream : public Stream {
    using CanonicalCode = Brotli::CanonicalCode;

public:
    enum class State {
        WindowSize,
        Idle,
        UncompressedData,
        CompressedCommand,
        CompressedLiteral,
        CompressedDistance,
        CompressedCopy,
        CompressedDictionary,
    };

    struct Block {
        size_t type;
        size_t type_previous;
        size_t number_of_types;

        size_t length;

        CanonicalCode type_code;
        CanonicalCode length_code;
    };

    class LookbackBuffer {
    private:
        LookbackBuffer(FixedArray<u8>& buffer)
            : m_buffer(move(buffer))
        {
        }

    public:
        static ErrorOr<LookbackBuffer> try_create(size_t size)
        {
            auto buffer = TRY(FixedArray<u8>::create(size));
            return LookbackBuffer { buffer };
        }

        void write(u8 value)
        {
            m_buffer[m_offset] = value;
            m_offset = (m_offset + 1) % m_buffer.size();
            m_total_written++;
        }

        u8 lookback(size_t offset) const
        {
            VERIFY(offset <= m_total_written);
            VERIFY(offset <= m_buffer.size());
            size_t index = (m_offset + m_buffer.size() - offset) % m_buffer.size();
            return m_buffer[index];
        }

        u8 lookback(size_t offset, u8 fallback) const
        {
            if (offset > m_total_written || offset > m_buffer.size())
                return fallback;
            VERIFY(offset <= m_total_written);
            VERIFY(offset <= m_buffer.size());
            size_t index = (m_offset + m_buffer.size() - offset) % m_buffer.size();
            return m_buffer[index];
        }

        size_t total_written() { return m_total_written; }

    private:
        FixedArray<u8> m_buffer;
        size_t m_offset { 0 };
        size_t m_total_written { 0 };
    };

public:
    BrotliDecompressionStream(MaybeOwned<Stream>);

    ErrorOr<Bytes> read_some(Bytes output_buffer) override;
    ErrorOr<size_t> write_some(ReadonlyBytes bytes) override { return m_input_stream.write_some(bytes); }
    bool is_eof() const override;
    bool is_open() const override { return m_input_stream.is_open(); }
    void close() override { m_input_stream.close(); }

private:
    ErrorOr<size_t> read_window_length();
    ErrorOr<size_t> read_size_number_of_nibbles();
    ErrorOr<size_t> read_variable_length();

    ErrorOr<void> read_context_map(size_t number_of_codes, Vector<u8>& context_map, size_t context_map_size);
    ErrorOr<void> read_block_configuration(Block&);

    ErrorOr<void> block_update_length(Block&);
    ErrorOr<void> block_read_new_state(Block&);

    size_t literal_code_index_from_context();

    LittleEndianInputBitStream m_input_stream;
    State m_current_state { State::WindowSize };
    Optional<LookbackBuffer> m_lookback_buffer;

    size_t m_window_size { 0 };
    bool m_read_final_block { false };
    size_t m_postfix_bits { 0 };
    size_t m_direct_distances { 0 };
    size_t m_distances[4] { 4, 11, 15, 16 };

    size_t m_bytes_left { 0 };
    size_t m_insert_length { 0 };
    size_t m_copy_length { 0 };
    bool m_implicit_zero_distance { false };
    size_t m_distance { 0 };
    ByteBuffer m_dictionary_data;

    Block m_literal_block;
    Vector<u8> m_literal_context_modes;
    Block m_insert_and_copy_block;
    Block m_distance_block;

    Vector<u8> m_context_mapping_literal;
    Vector<u8> m_context_mapping_distance;

    Vector<CanonicalCode> m_literal_codes;
    Vector<CanonicalCode> m_insert_and_copy_codes;
    Vector<CanonicalCode> m_distance_codes;
};

}