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

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

#include <util/string/join.h>

namespace NH2Log {
    namespace {
        struct TRawHeaderBlock {
            ui32 StreamId = 0;
            TDeque<TString> Blocks;
            bool Closed = false;
        };

        const TStringBuf RawHeaderSessionMarker{"RawHeaderSession"};
        const TStringBuf RawHeaderBlockMarker{"HeaderBlock"};

        struct TCompressionErrorSession : public TSession {
            TDeque<TRawHeaderBlock> HeaderBlocks;
            bool HasStart = false;
        };
    }

    bool IsHeaderSessionStart(const TLogEntry& item) {
        return item.Event.EventType == EEventType::Entry && item.Event.Name == RawHeaderSessionMarker;
    }

    bool IsHeaderSessionEnd(const TLogEntry& item) {
        return item.Event.EventType == EEventType::Exit && item.Event.Name == RawHeaderSessionMarker;
    }

    bool IsRecvHeaderBlock(const TLogEntry& item) {
        return item.Event.EventType == EEventType::Event && item.Event.Name == RawHeaderBlockMarker;
    }

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

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

            if (IsH2SessionStart(item)) {
                sess->HasStart = true;
            }

            if (sess->HasStart) {
                if (IsRecvHeadersBlockStart(item)) {
                    sess->HeaderBlocks.emplace_back().StreamId = *ParseStreamId(item);
                }

                if (IsRecvHeadersPayload(item) || IsRecvContinuationPayload(item)) {
                    Y_ENSURE(!sess->HeaderBlocks.empty());
                    Y_ENSURE(!sess->HeaderBlocks.back().Closed);
                    sess->HeaderBlocks.back().Blocks.emplace_back(ExtractRecvHeaderBlock(item));
                }

                if (IsRecvHeadersBlockEnd(item)) {
                    Y_ENSURE(!sess->HeaderBlocks.empty());
                    Y_ENSURE(!sess->HeaderBlocks.back().Closed);
                    sess->HeaderBlocks.back().Closed = true;
                }
            }

            if (IsH2SessionEnd(item)) {
                if (sess->HasStart) {
                    auto& out = fileManager.GetFileOutput(fileManager.GetFileNameBySuffix("recv_headers"));
                    out << item.RawPrefix << "\t" << RawHeaderSessionMarker << "\t\n";
                    for (const auto& hb : sess->HeaderBlocks) {
                        TString res;
                        for (const auto& b : hb.Blocks) {
                            res += Base64Decode(b);
                        }
                        if (hb.StreamId && hb.Closed) {
                            out << item.RawPrefix << "\t[Event] " << RawHeaderBlockMarker << "\t{StreamId="
                                << hb.StreamId << ", EscBase64(Block)=" << Base64Encode(res) << "}\n";
                        } else {
                            out << item.RawPrefix << "\t[Event] PartialHeaderBlock\t{StreamId="
                                << hb.StreamId << ", len=" << res.size() << "}\n";
                            break;
                        }
                    }
                    out << item.RawPrefix << "\t" << RawHeaderSessionMarker << " [Done]\t\n";
                }
                active.erase(item.Addr);
            }
        }
    }
}
