#include "hpack_integers.h"
#include "hpack_strings.h"
#include "hpack_writer.h"

#include <util/generic/strbuf.h>

namespace NSrvKernel::NHTTP2 {
    THPackWriter::THPackWriter(const THPackEncoderSettings& settings) noexcept
        : Settings_(settings)
        , ChunkBuffer_(Settings_.MaxFrameSize)
        , ChunkOutput_(ChunkBuffer_)
    {
        Y_VERIFY(settings.MaxFrameSize > 0);
        NextChunk(Settings_.MaxFrameSize);
    }

    TChunkList& THPackWriter::GetChunkList() noexcept {
        return ChunkList_;
    }

    void THPackWriter::Commit() noexcept {
        NextChunk(0);
    }

    void THPackWriter::EnsureBuffer(size_t needed) noexcept {
        if (ChunkOutput_.SizeAvailable() < needed) {
            NextChunk(std::max<size_t>(needed, Settings_.MaxFrameSize));
        }
    }

    void THPackWriter::NextChunk(size_t nextSize) noexcept {
        if (ChunkOutput_.SizeConsumed()) {
            ChunkBuffer_.Resize(ChunkOutput_.SizeConsumed());
            ChunkList_.Push(NewChunk(std::move(ChunkBuffer_)));
        }
        ChunkBuffer_.Resize(nextSize);
        ChunkOutput_ = TOutputRegion(ChunkBuffer_);
    }

    void THPackWriter::WriteHeaderName(TStringBuf name) noexcept {
        WriteString(name, SM_NAME);
    }

    void THPackWriter::WriteHeaderValue(TStringBuf value, bool allowHuffman) noexcept {
        WriteString(value, allowHuffman ? SM_VALUE : SM_VALUE_NO_HUFFMAN);
    }

    void THPackWriter::WriteString(TStringBuf input, EStringMode mode) noexcept {
        if (Settings_.EnableHuffman && SM_VALUE_NO_HUFFMAN != mode) {
            const auto huffmanSize = GetHuffmanEncodedSize(input, mode);

            if (huffmanSize < input.size()) {
                DoWriteStringHuffman(input, huffmanSize, mode);
            } else {
                DoWriteStringRaw(input, mode);
            }
        } else {
            DoWriteStringRaw(input, mode);
        }
    }

    void THPackWriter::DoWriteStringRaw(TStringBuf input, EStringMode mode) noexcept {
        const auto sizeSize = GetEncodedIntSize<FF_STRING_HUFFMAN_MASK>(input.size());
        EnsureBuffer(input.size() + sizeSize);
        WriteInt<FF_STRING_HUFFMAN_MASK, FF_STRING_HUFFMAN_OFF_FLAG>(input.size(), ChunkOutput_);

        if (SM_NAME == mode) {
            CopyStringLowerCase(input, ChunkOutput_);
        } else {
            CopyString(input, ChunkOutput_);
        }
    }

    void THPackWriter::DoWriteStringHuffman(TStringBuf input, size_t encodedSize, EStringMode mode) noexcept {
        const auto sizeSize = GetEncodedIntSize<FF_STRING_HUFFMAN_MASK>(encodedSize);
        EnsureBuffer(encodedSize + sizeSize);
        WriteInt<FF_STRING_HUFFMAN_MASK, FF_STRING_HUFFMAN_ON_FLAG>(encodedSize, ChunkOutput_);

        if (SM_NAME == mode) {
            HuffmanEncodeStringLowerCase(input, ChunkOutput_);
        } else {
            HuffmanEncodeString(input, ChunkOutput_);
        }
    }

    size_t THPackWriter::GetHuffmanEncodedSize(TStringBuf input, EStringMode mode) noexcept {
        if (SM_NAME == mode) {
            return GetHuffmanEncodedLowerCaseStringSize(input);
        } else {
            return GetHuffmanEncodedStringSize(input);
        }
    }

    void THPackWriter::WriteTableSizeUpdate(ui32 newSize) noexcept {
        EnsureBuffer(GetEncodedIntSize<FF_TABLE_SIZE_UPDATE_MASK>(newSize));
        WriteInt<FF_TABLE_SIZE_UPDATE_MASK, FF_TABLE_SIZE_UPDATE_FLAG>(newSize, ChunkOutput_);
    }

    void THPackWriter::WriteIndexedAll(ui32 index) noexcept {
        EnsureBuffer(GetEncodedIntSize<FF_HEADER_FIELD_INDEXED_ALL_MASK>(index));
        WriteInt<FF_HEADER_FIELD_INDEXED_ALL_MASK, FF_HEADER_FIELD_INDEXED_ALL_FLAG>(index, ChunkOutput_);
    }

    void THPackWriter::WriteHeaderField(EHPackHeaderRule rule, ui32 nameIndex) noexcept {
        switch (rule) {
        case EHPackHeaderRule::Index:
            EnsureBuffer(GetEncodedIntSize<FF_HEADER_FIELD_INDEX_MASK>(nameIndex));
            WriteInt<FF_HEADER_FIELD_INDEX_MASK, FF_HEADER_FIELD_INDEX_FLAG>(nameIndex, ChunkOutput_);
            break;
        case EHPackHeaderRule::DoNotIndex:
            EnsureBuffer(GetEncodedIntSize<FF_HEADER_FIELD_NO_INDEX_MASK>(nameIndex));
            WriteInt<FF_HEADER_FIELD_NO_INDEX_MASK, FF_HEADER_FIELD_DO_NOT_INDEX_FLAG>(nameIndex, ChunkOutput_);
            break;
        case EHPackHeaderRule::NeverIndex:
            EnsureBuffer(GetEncodedIntSize<FF_HEADER_FIELD_NO_INDEX_MASK>(nameIndex));
            WriteInt<FF_HEADER_FIELD_NO_INDEX_MASK, FF_HEADER_FIELD_NEVER_INDEX_FLAG>(nameIndex, ChunkOutput_);
            break;
        default:
            Y_FAIL("Invalid enum");
        }
    }

}
