#include "hpack_table.h"
#include "hpack_static_table.h"

#include <library/cpp/scheme/scheme.h>

#include <util/generic/strbuf.h>

namespace NSrvKernel::NHTTP2 {
    namespace {
        [[nodiscard]]
        inline size_t DoGetHeaderFieldSize(THeaderFieldView headerField) noexcept {
            return headerField.Name.size() + headerField.Value.size();
        }

        [[nodiscard]]
        inline size_t DoGetHeaderFieldHPackSize(THeaderFieldView headerField) noexcept {
            return DoGetHeaderFieldSize(headerField) + RFC_HEADER_FIELD_OVERHEAD;
        }

        [[nodiscard]]
        inline TChunkPtr DoCopyString(TStringBuf str) noexcept {
            return NewChunk(str.size(), NewChunkData(str));
        }

        [[nodiscard]]
        inline THeaderField DoCopyHeaderField(THeaderFieldView headerField) noexcept {
            return {DoCopyString(headerField.Name), DoCopyString(headerField.Value)};
        }

        [[nodiscard]]
        inline NImpl::THeaderFieldBlobPtr DoMakeHeaderFieldBlob(THeaderFieldView headerField) noexcept {
            return new (DoGetHeaderFieldSize(headerField)) NImpl::THeaderFieldBlob(headerField);
        }
    }

    namespace NImpl {
        THeaderFieldBlob::THeaderFieldBlob(THeaderFieldView headerField) noexcept
            : NameSize_(headerField.Name.size())
            , HasValue_(headerField.HasValue)
            , ValueSize_(headerField.Value.size())
        {
            Y_VERIFY(DoGetHeaderFieldSize(headerField) <= AdditionalLength());

            if (NameSize_) {
                memcpy(
                    (char*) AdditionalData(),
                    headerField.Name.data(),
                    NameSize_
                );
            }
            if (ValueSize_) {
                memcpy(
                    (char*) AdditionalData() + NameSize_,
                    headerField.Value.data(),
                    ValueSize_
                );
            }
        }

        [[nodiscard]]
        THeaderFieldView THeaderFieldBlob::GetHeaderField() const noexcept {
            TStringBuf name{(const char*)AdditionalData(), NameSize_};
            TStringBuf value{(const char*)AdditionalData() + NameSize_, ValueSize_};
            return HasValue_ ? THeaderFieldView{name, value} : THeaderFieldView{name};
        }
    }

    THPackTableBase::THPackTableBase(const THPackBaseSettings& settigs) noexcept
        : MaxSize_(settigs.MaxHeaderTableSize)
    {}

    void THPackTableBase::SetMaxSize(size_t newSize) noexcept {
        if (newSize < MaxSize_) {
            DoUpdateSize(newSize);
        }
        MaxSize_ = newSize;
    }

    void THPackTableBase::DoUpdateSize(ui32 newSize) noexcept {
        const auto oldSize = TableSize_;
        TableSize_ = newSize;
        if (TableSize_ < oldSize) {
            DoShrinkTable();
        }
    }

    void THPackTableBase::DoShrinkTable() noexcept {
        while (HeaderFieldsTable_ && CurrentSize_ > TableSize_) {
            auto headerFieldPtr = HeaderFieldsTable_.back();
            Y_VERIFY(headerFieldPtr);

            const auto headerField = headerFieldPtr->GetHeaderField();
            const auto headerFieldSize = DoGetHeaderFieldHPackSize(headerField);

            HeaderFieldsTable_.pop_back();
            CurrentSize_ -= headerFieldSize;

            DoRemoveHeaderField(headerField);
        }

        const size_t tableCount = HeaderFieldsTable_.size();
        Y_VERIFY(CurrentSize_ >= tableCount * RFC_HEADER_FIELD_OVERHEAD);
    }

    void THPackTableBase::DumpDynamicTable(size_t& sz, THeadersList& res) const noexcept {
        sz = CurrentSize_;
        for (const auto& headerFieldPtr : HeaderFieldsTable_) {
            auto headerField = headerFieldPtr->GetHeaderField();
            res.Add(TChunkList{NewChunkNonOwning(headerField.Name)},
                    TChunkList{NewChunkNonOwning(headerField.Value)});
        }
    }

    void THPackTableBase::AddHeaderField(THeaderFieldView headerField) noexcept {
        // It is ok to add headers before cleaning the table since we limit the maximum header size.
        CurrentSize_ += DoGetHeaderFieldHPackSize(headerField);
        HeaderFieldsTable_.push_front(DoGetHeaderFieldBlob(headerField));
        DoShrinkTable();
    }

    NImpl::THeaderFieldBlobPtr THPackTableBase::DoGetHeaderFieldBlob(
        THeaderFieldView headerField) noexcept
    {
        return DoMakeHeaderFieldBlob(headerField);
    }

    TError THPackDecoderTable::UpdateSize(ui32 newSize) noexcept {
        Y_REQUIRE(newSize <= MaxSize_,
            CompressionError(ECompressionError::InvalidSizeUpdate));
        DoUpdateSize(newSize);
        return {};
    }

    TErrorOr<THeaderFieldView> THPackDecoderTable::DoGetHeaderField(ui32 index) const noexcept {
        Y_VERIFY(index > 0);

        if (index <= RFC_STATIC_TABLE_INDEX_MAX) {
            return TStaticTable::Instance().GetHeaderField(index);
        } else {
            index -= RFC_STATIC_TABLE_INDEX_MAX + 1;

            Y_REQUIRE(index < HeaderFieldsTable_.size(),
                CompressionError(ECompressionError::InvalidIndex) << "table: " << HeaderFieldsTable_.size() << " index: " << index);
            const auto& headerFieldPtr = HeaderFieldsTable_.at(index);
            Y_VERIFY(headerFieldPtr);
            return headerFieldPtr->GetHeaderField();
        }
    }

    TErrorOr<THeaderField> THPackDecoderTable::GetHeaderField(ui32 index, THPackLimits& limits) const noexcept {
        THeaderFieldView fieldView;
        Y_PROPAGATE_ERROR(DoGetHeaderField(index).AssignTo(fieldView));
        Y_PROPAGATE_ERROR(limits.AcceptLength(DoGetHeaderFieldSize(fieldView)));
        return DoCopyHeaderField(fieldView);
    }

    TErrorOr<TChunkPtr> THPackDecoderTable::GetHeaderName(ui32 index, THPackLimits& limits) const noexcept {
        THeaderFieldView headerFieldView;
        Y_PROPAGATE_ERROR(DoGetHeaderField(index).AssignTo(headerFieldView));
        Y_PROPAGATE_ERROR(limits.AcceptLength(headerFieldView.Name.size()));
        return DoCopyString(headerFieldView.Name);
    }

    ui32 THPackEncoderTable::GetHeaderFieldIndex(THeaderFieldView headerField) const noexcept {
        if (ui32 staticIdx = TStaticTable::Instance().GetHeaderFieldIndex(headerField)) {
            return staticIdx;
        } else if (const auto* ptr = HeaderFieldsIndex_.FindPtr(headerField)) {
            // NOTE: here we use Epoch, not Epoch + 1 because Epoch is already +1
            const ui64 idx = Epoch_ - ptr->Epoch_ + RFC_STATIC_TABLE_INDEX_MAX;
            Y_VERIFY(idx <= IMPL_HPACK_HEADER_TABLE_INDEX_LIMIT);
            return (ui32) idx;
        } else {
            return 0;
        }
    }

    ui32 THPackEncoderTable::GetHeaderNameIndex(TStringBuf name) const noexcept {
        return GetHeaderFieldIndex(name);
    }

    NImpl::THeaderFieldBlobPtr THPackEncoderTable::DoInsertIndex(THeaderFieldView headerField) noexcept {
        NImpl::THeaderFieldsIndex::insert_ctx ins;
        auto it = HeaderFieldsIndex_.find(headerField, ins);
        if (it == HeaderFieldsIndex_.end()) {
            NImpl::THeaderFieldBlobPtr ptr = DoMakeHeaderFieldBlob(headerField);
            HeaderFieldsIndex_.insert_direct(NImpl::THeaderFieldsIndex::value_type{ptr, NImpl::TIndexRef(Epoch_)}, ins);
            return ptr;
        } else {
            it->second.Epoch_ = Epoch_;
            it->second.RefCount_ += 1;
            return it->first;
        }
    }

    void THPackEncoderTable::DoRemoveIndex(THeaderFieldView headerField) noexcept {
        auto it = HeaderFieldsIndex_.find(headerField);
        Y_VERIFY(it != HeaderFieldsIndex_.end());
        it->second.RefCount_ -= 1;
        if (!it->second.RefCount_) {
            HeaderFieldsIndex_.erase(it);
        }
    }

    NImpl::THeaderFieldBlobPtr THPackEncoderTable::DoGetHeaderFieldBlob(
        THeaderFieldView headerField) noexcept
    {
        auto ptr = DoInsertIndex(headerField);
        auto namePtr = DoInsertIndex(headerField.Name);
        Y_UNUSED(namePtr);
        Epoch_ += 1;
        return ptr;
    }

    void THPackEncoderTable::DoRemoveHeaderField(THeaderFieldView headerField) noexcept {
        DoRemoveIndex(headerField);
        DoRemoveIndex(headerField.Name);
    }
}
