#include <util/string/hex.h>
#include "hpack_integers.h"
#include "hpack_reader.h"
#include "hpack_strings.h"

namespace NSrvKernel::NHTTP2 {
    namespace {
        class TChunkDataBuffer {
        public:
            explicit TChunkDataBuffer(size_t maxSize) noexcept
                : ChunkDataPtr_(NewChunkData(maxSize))
                , OutputRegion_(*ChunkDataPtr_)
            {}

            [[nodiscard]] TOutputRegion& GetOutputRegion() noexcept {
                return OutputRegion_;
            }

            [[nodiscard]] size_t GetTotalConsumed() const noexcept {
                return OutputRegion_.SizeConsumed();
            }

            [[nodiscard]] TChunkPtr FinishChunk() noexcept {
                TChunkPtr res = NewChunk(OutputRegion_.SizeConsumed(), ChunkDataPtr_.Release());
                OutputRegion_.Clear();
                return res;
            }

        private:
            THolder<TChunkData> ChunkDataPtr_;
            TOutputRegion OutputRegion_;
        };
    }

    THPackReader::THPackReader(TInputRegion headerBlock) noexcept
        : HeaderBlock_(headerBlock)
    {}

    bool THPackReader::HasTableSizeUpdate() const noexcept {
        return HeaderBlock_.SizeAvailable() && ((HeaderBlock_[0] & FF_TABLE_SIZE_UPDATE_MASK) == FF_TABLE_SIZE_UPDATE_FLAG);
    }

    TErrorOr<ui32> THPackReader::ReadTableSizeUpdate() noexcept {
        return ReadInt<FF_TABLE_SIZE_UPDATE_MASK>(HeaderBlock_);
    }

    TErrorOr<EHeaderFieldType> THPackReader::PeekNextHeaderFieldType() const noexcept {
        Y_VERIFY(HeaderBlock_.SizeAvailable());

        const ui8 next = HeaderBlock_[0];

        if ((next & FF_HEADER_FIELD_INDEXED_ALL_MASK) == FF_HEADER_FIELD_INDEXED_ALL_FLAG) {
            return EHeaderFieldType::IndexedAll;
        } else if ((next & FF_HEADER_FIELD_INDEX_MASK) == FF_HEADER_FIELD_INDEX_FLAG) {
            return EHeaderFieldType::ToIndex;
        } else if ((next & FF_HEADER_FIELD_NO_INDEX_MASK) == FF_HEADER_FIELD_DO_NOT_INDEX_FLAG) {
            return EHeaderFieldType::NotToIndex;
        } else if ((next & FF_HEADER_FIELD_NO_INDEX_MASK) == FF_HEADER_FIELD_NEVER_INDEX_FLAG) {
            return EHeaderFieldType::NeverIndex;
        } else {
            return Y_MAKE_ERROR(CompressionError(ECompressionError::InvalidHeaderFieldType));
        }
    }

    TErrorOr<ui32> THPackReader::ReadNextHeaderFieldIndex(EHeaderFieldType hft) noexcept {
        switch (hft) {
        case EHeaderFieldType::IndexedAll:
            return ReadInt<FF_HEADER_FIELD_INDEXED_ALL_MASK>(HeaderBlock_);
        case EHeaderFieldType::ToIndex:
            return ReadInt<FF_HEADER_FIELD_INDEX_MASK>(HeaderBlock_);
        case EHeaderFieldType::NotToIndex:
        case EHeaderFieldType::NeverIndex:
            return ReadInt<FF_HEADER_FIELD_NO_INDEX_MASK>(HeaderBlock_);
        default:
            Y_FAIL("Invalid enum");
        }
    }

    // RFC 7540: A server that receives a larger header block than it is willing to handle can send an HTTP 431
    //           (Request Header Fields Too Large) status code [RFC6585]. A client can discard responses
    //           that it cannot process. The header block MUST be processed to ensure a consistent connection state,
    //           unless the connection is closed.
    //
    // Impl: So we chose to close the whole connection for any client sending too big headers.
    //       It's both easier and corresponds to what we do in http.
    TErrorOr<TChunkPtr> THPackReader::ReadString(THPackLimits& limits) noexcept {
        Y_REQUIRE(HeaderBlock_.SizeAvailable(),
            CompressionError(ECompressionError::HeaderBlockEnd));

        const auto prefix = HeaderBlock_[0];
        const bool huffman = ((prefix & FF_STRING_HUFFMAN_MASK) == FF_STRING_HUFFMAN_ON_FLAG);

        ui32 length = 0;
        Y_PROPAGATE_ERROR(ReadInt<FF_STRING_HUFFMAN_MASK>(HeaderBlock_).AssignTo(length));

        Y_REQUIRE(length <= HeaderBlock_.SizeAvailable(),
            CompressionError(ECompressionError::HeaderBlockEnd)
                << "expected " << length << ", got " << HeaderBlock_.SizeAvailable());

        if (huffman) {
            return DoReadStringHuffman(limits, length);
        } else {
            return DoReadStringRaw(limits, length);
        }
    }

    TErrorOr<TChunkPtr> THPackReader::DoReadStringRaw(THPackLimits& limits, ui32 length) noexcept {
        TChunkDataBuffer buffer{length};

        Y_PROPAGATE_ERROR(limits.AcceptLength(length));
        CopyString(HeaderBlock_.AsStringBuf(length), buffer.GetOutputRegion());
        HeaderBlock_.Consume(length);

        return buffer.FinishChunk();
    }

    TErrorOr<TChunkPtr> THPackReader::DoReadStringHuffman(THPackLimits& limits, ui32 length) noexcept {
        TChunkDataBuffer buffer{std::min(GetMaxHuffmanDecodedStringSize(length),
                                limits.GetAllowedLength())};

        Y_PROPAGATE_ERROR(HuffmanDecodeString(HeaderBlock_.AsStringBuf(length), buffer.GetOutputRegion()));
        limits.AddLength(buffer.GetTotalConsumed());
        HeaderBlock_.Consume(length);

        return buffer.FinishChunk();
    }
}
