#include "hpack.h"
#include "hpack_reader.h"
#include "hpack_writer.h"
#include "hpack_table.h"

#include <util/generic/maybe.h>

#include <algorithm>

namespace NSrvKernel::NHTTP2 {
    [[nodiscard]]
    size_t GetMinHeaderListSize(size_t headerBlockSize) noexcept {
        return GetMinHuffmanDecodedStringSize(headerBlockSize);
    }


    // THPackBase ======================================================================================================

    void THPackBase::DumpDynamicTable(size_t& sz, THeadersList& res) noexcept {
        GetTable().DumpDynamicTable(sz, res);
    }


    // THPackDecoder ===================================================================================================

    THPackDecoder::THPackDecoder(const THPackDecoderSettings& settings) noexcept
        : Settings_(settings)
        , Table_(settings)
    {}

    THPackTableBase& THPackDecoder::GetTable() noexcept {
        return Table_;
    }

    TErrorOr<THeadersList> THPackDecoder::Decode(TStringBuf headerBlock) noexcept {
        THeadersList headers;

        THPackReader reader{headerBlock};
        THPackLimits limits{Settings_};

        {
            TMaybe<ui32> pendingSizeUpdate;
            TMaybe<ui32> minimalSizeUpdate;

            while (reader.HasTableSizeUpdate()) {
                ui32 newSize = 0;
                Y_PROPAGATE_ERROR(reader.ReadTableSizeUpdate().AssignTo(newSize));
                minimalSizeUpdate = minimalSizeUpdate ? std::min(*minimalSizeUpdate, newSize) : newSize;
                pendingSizeUpdate = newSize;
            }

            if (pendingSizeUpdate) {
                if (minimalSizeUpdate && *minimalSizeUpdate != *pendingSizeUpdate) {
                    Y_PROPAGATE_ERROR(Table_.UpdateSize(*minimalSizeUpdate));
                }

                Y_PROPAGATE_ERROR(Table_.UpdateSize(*pendingSizeUpdate));
            }
        }

        while (reader.HasNext()) {
            Y_PROPAGATE_ERROR(limits.AcceptHeader());

            EHeaderFieldType hfType = EHeaderFieldType::NeverIndex;
            Y_PROPAGATE_ERROR(reader.PeekNextHeaderFieldType().AssignTo(hfType));
            ui32 index = 0;
            Y_PROPAGATE_ERROR(reader.ReadNextHeaderFieldIndex(hfType).AssignTo(index));

            THeaderField headerField;

            if (EHeaderFieldType::IndexedAll == hfType) {
                Y_REQUIRE(index,
                    CompressionError(ECompressionError::InvalidIndex));
                Y_PROPAGATE_ERROR(Table_.GetHeaderField(index, limits).AssignTo(headerField));
            } else {
                if (index) {
                    Y_PROPAGATE_ERROR(Table_.GetHeaderName(index, limits).AssignTo(headerField.first));
                } else {
                    Y_PROPAGATE_ERROR(reader.ReadString(limits).AssignTo(headerField.first));
                }
                Y_PROPAGATE_ERROR(reader.ReadString(limits).AssignTo(headerField.second));
            }

            Y_VERIFY(headerField.first);
            Y_VERIFY(headerField.second);

            THolder<THeader> header{new THeader};

            switch (hfType) {
            case EHeaderFieldType::NotToIndex:
                header->HPackHeaderRule = EHPackHeaderRule::DoNotIndex;
                break;
            case EHeaderFieldType::NeverIndex:
                header->HPackHeaderRule = EHPackHeaderRule::NeverIndex;
                break;
            case EHeaderFieldType::ToIndex:
                Table_.AddHeaderField(
                    {headerField.first->AsStringBuf(), headerField.second->AsStringBuf()}
                );
                [[fallthrough]];
            case EHeaderFieldType::IndexedAll:
                header->HPackHeaderRule = EHPackHeaderRule::Index;
                break;
            }

            header->Key = TChunkList{std::move(headerField.first)};
            header->Value = TChunkList{std::move(headerField.second)};

            headers.PushBack(header.Release());
        }

        return headers;
    }

    void THPackDecoder::UpdateSettings(const THPackDecoderSettings& settings) {
        if (settings.MaxHeaderTableSize != Settings_.MaxHeaderTableSize) {
            // The client code is expected to validate the newSize against a sane limit before submitting.
            // Note that IMPL_HEADER_TABLE_SIZE_LIMIT is way too big.
            const auto newSize = settings.MaxHeaderTableSize;
            Y_VERIFY(newSize <= IMPL_HEADER_TABLE_SIZE_LIMIT);
            GetTable().SetMaxSize(newSize);
        }

        Settings_ = settings;
    }


    // THPackEncoder ===================================================================================================

    THPackEncoder::THPackEncoder(const THPackEncoderSettings& settings) noexcept
        : Settings_(settings)
        , Table_(settings)
    {}

    THPackTableBase& THPackEncoder::GetTable() noexcept {
        return Table_;
    }

    TChunkList THPackEncoder::Encode(const THeadersList& headers) noexcept {
        THPackWriter writer{Settings_};

        if (PendingSizeUpdate_) {
            if (MinimalSizeUpdate_ && *MinimalSizeUpdate_ != *PendingSizeUpdate_) {
                writer.WriteTableSizeUpdate(*MinimalSizeUpdate_);
                MinimalSizeUpdate_.Clear();
            }

            writer.WriteTableSizeUpdate(*PendingSizeUpdate_);
            PendingSizeUpdate_.Clear();
        }

        for (const auto& header : headers) {
            TUnitedChunkList headerName{header.Key};
            TUnitedChunkList headerValue{header.Value};

            auto index = Table_.GetHeaderFieldIndex(
                {headerName.AsStringBuf(), headerValue.AsStringBuf()}
            );

            auto headerRule = header.HPackHeaderRule.GetOrElse(
                index <= RFC_STATIC_TABLE_INDEX_MAX ? EHPackHeaderRule::Index : Settings_.DefaultHeaderRule
            );

            if (index && EHPackHeaderRule::NeverIndex != headerRule) {
                writer.WriteIndexedAll(index);
            } else {
                index = Table_.GetHeaderNameIndex(headerName.AsStringBuf());
                writer.WriteHeaderField(headerRule, index);

                if (!index) {
                    writer.WriteHeaderName(headerName.AsStringBuf());
                }

                writer.WriteHeaderValue(
                    headerValue.AsStringBuf(),
                    EHPackHeaderRule::NeverIndex != headerRule
                );

                if (EHPackHeaderRule::Index == headerRule) {
                    Table_.AddHeaderField(
                        {headerName.AsStringBuf(), headerValue.AsStringBuf()}
                    );
                }
            }
        }

        writer.Commit();

        return std::move(writer.GetChunkList());
    }

    void THPackEncoder::UpdateSettings(const THPackEncoderSettings& settings) {
        UpdateHeaderTableSize(settings.MaxHeaderTableSize);
        Settings_ = settings;
    }

    void THPackEncoder::UpdateHeaderTableSize(ui32 newSize) {
        // The client code is expected to validate the newSize against a sane limit before submitting.
        // Note that IMPL_HEADER_TABLE_SIZE_LIMIT is way too big.
        Y_VERIFY(newSize <= IMPL_HEADER_TABLE_SIZE_LIMIT);
        MinimalSizeUpdate_ = MinimalSizeUpdate_ ? std::min(*MinimalSizeUpdate_, newSize) : newSize;
        PendingSizeUpdate_ = newSize;
        GetTable().SetMaxSize(newSize);
        Settings_.MaxHeaderTableSize = newSize;
    }
}
