#pragma once
#include <util/generic/vector.h>
#include <util/generic/map.h>
#include <util/generic/string.h>
#include <library/cpp/json/writer/json_value.h>
#include <library/cpp/logger/global/global.h>
#include <util/system/rwlock.h>

class TMessagesCollector {
private:
    TMap<TString, TVector<TString>> Messages;

public:

    Y_FORCE_INLINE bool HasMessages() const {
        return Messages.size();
    }

    void MergeMessages(const TMessagesCollector& collector) {
        for (auto&& i : collector.Messages) {
            for (auto&& j : i.second) {
                Messages[i.first].push_back(j);
            }
        }
    }

    template <ELogPriority LogLevel = TLOG_ERR>
    void AddMessage(const TString& source, const TString& message, bool doThrowException = false) {
        Messages[source].push_back(message);
        TEMPLATE_LOG(LogLevel) << source << " : " << message << Endl;
        if (doThrowException) {
            ythrow yexception() << message;
        }
    }

    template <ELogPriority LogLevel = TLOG_ERR>
    void AddMessage(const TSourceLocation& sl, const TString& message, bool doThrowException = false) {
        Messages[ToString(sl)].push_back(message);
        TEMPLATE_LOG(LogLevel) << "(" << sl << ")" << " : " << message << Endl;
        if (doThrowException) {
            ythrow yexception() << sl << ": " << message;
        }
    }

    void ClearMessages() {
        Messages.clear();
    }

    TString GetStringReport() const {
        return GetReport().GetStringRobust();
    }

    NJson::TJsonValue GetReport() const {
        NJson::TJsonValue result(NJson::JSON_MAP);
        for (auto&& i : Messages) {
            for (auto&& m : i.second) {
                result[i.first].AppendValue(m);
            }
        }
        return result;
    }
};

class TMTMessagesCollector: public TMessagesCollector {
private:
    TRWMutex Mutex;
public:
    template <ELogPriority LogLevel = TLOG_ERR>
    void AddMessage(const TString& source, const TString& message, bool doThrowException = false) {
        TWriteGuard wg(Mutex);
        TMessagesCollector::AddMessage<LogLevel>(source, message, doThrowException);
    }

    template <ELogPriority LogLevel = TLOG_ERR>
    void AddMessage(const TSourceLocation& sl, const TString& message, bool doThrowException = false) {
        TWriteGuard wg(Mutex);
        TMessagesCollector::AddMessage<LogLevel>(sl, message, doThrowException);
    }

    void Clear() {
        TWriteGuard wg(Mutex);
        TMessagesCollector::ClearMessages();
    }

    TString GetStringReport() const {
        return GetReport().GetStringRobust();
    }

    NJson::TJsonValue GetReport() const {
        TReadGuard wg(Mutex);
        return TMessagesCollector::GetReport();
    }
};
