#include "aux_utils.h"
#include "log_parsers.h"
#include "parse_raw_header_blocks.h"
#include "dump_raw_header_blocks.h"

#include <balancer/kernel/http2/server/hpack/hpack_definitions.h>
#include <balancer/kernel/http2/server/hpack/hpack.h>

#include <library/cpp/string_utils/base64/base64.h>

#include <util/generic/vector.h>

namespace NH2Log {
    namespace {
        using namespace NSrvKernel;
        using namespace NSrvKernel::NHTTP2;

        const TStringBuf HeaderSessionMarker{"HeaderSession"};
        const TStringBuf HeaderMarker{"Header"};
        const TStringBuf HeaderBlockSizeMarker{"HeaderBlockSize"};

        struct TRawHeaderBlock {
            ui32 StreamId = 0;
            TString Block;
            THeadersList Headers;
        };

        struct THeadersErrorSession : public TSession {
            static THPackDecoderSettings MakeSettings() {
                THPackDecoderSettings res;
                res.MaxHeaderTableSize = IMPL_HEADER_TABLE_SIZE_LIMIT;
                res.MaxHeadersCount = -1;
                res.MaxHeaderListSize = -1;
                return res;
            }

            THPackDecoder Decoder{MakeSettings()};
            TDeque<TRawHeaderBlock> HeaderBlocks;
            bool HasStart = false;
        };
    }

    void ParseRawClientHeaderBlocks(TFileManager& fileManager) {
        Y_H2_LOG_FUNC(fileManager);
        THashMap<TString, THeadersErrorSession> active;
        THeadersErrorSession* sess = nullptr;
        TString line;
        while (fileManager.GetInput().ReadLine(line)) {
            if (!(sess = ParseNextLogLine(sess, active, line))) {
                continue;
            }

            auto& item = sess->Log.back();

            if (IsHeaderSessionStart(item)) {
                sess->HasStart = true;
            } else if (sess->HasStart && IsRecvHeaderBlock(item)) {
                auto& block = sess->HeaderBlocks.emplace_back();
                block.StreamId = *ParseStreamId(item);
                block.Block = Base64Decode(ExtractRecvHeaderBlock(item));

                TryRethrowError(sess->Decoder.Decode(block.Block).AssignTo(block.Headers));
            } else if (IsHeaderSessionEnd(item)) {
                if (sess->HasStart) {
                    auto& out = fileManager.GetFileOutput(fileManager.GetFileNameBySuffix("parsed_headers"));
                    out << item.Addr << "\t" << HeaderSessionMarker << "\t\n";
                    for (const auto& hb : sess->HeaderBlocks) {
                        out << item.Addr << "\t" << HeaderBlockSizeMarker << "\tStreamId="
                            << hb.StreamId << "\tHeaderBlockSize=" << hb.Block.size() << "\n";

                        for (const auto& h : hb.Headers) {
                            out << item.Addr << "\t" << HeaderMarker << "\tStreamId="
                                << hb.StreamId << "\tHeaderSize="
                                << (h.Key.size() + h.Value.size()) << "\t" << TString{Union(h.Key)->AsStringBuf()}.Quote() << ":\t"
                                << TString{Union(h.Value)->AsStringBuf()}.Quote() << "\n";
                        }
                    }
                    out << item.Addr << "\t" << HeaderSessionMarker << " [Done]\t\n";
                }
                active.erase(item.Addr);
            }
        }
    }
}
