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

#pragma once

#include <AK/CircularQueue.h>
#include <AK/FixedArray.h>
#include <LibCore/InputBitStream.h>
#include <LibCore/Stream.h>

namespace Compress {

using Core::Stream::LittleEndianInputBitStream;
using Core::Stream::Stream;

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

    class CanonicalCode {
        friend class BrotliDecompressionStream;

    public:
        CanonicalCode() = default;
        ErrorOr<size_t> read_symbol(LittleEndianInputBitStream&);
        void clear()
        {
            m_symbol_codes.clear();
            m_symbol_values.clear();
        }

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

    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>::try_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(Stream&);

    bool is_readable() const override { return m_input_stream.is_readable(); }
    ErrorOr<Bytes> read(Bytes output_buffer) override;
    bool is_writable() const override { return m_input_stream.is_writable(); }
    ErrorOr<size_t> write(ReadonlyBytes bytes) override { return m_input_stream.write(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<size_t> read_complex_prefix_code_length();

    ErrorOr<void> read_prefix_code(CanonicalCode&, size_t alphabet_size);
    ErrorOr<void> read_simple_prefix_code(CanonicalCode&, size_t alphabet_size);
    ErrorOr<void> read_complex_prefix_code(CanonicalCode&, size_t alphabet_size, size_t hskip);
    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;
};

}