#pragma once

#include "sender_context.h"

#include <saas/api/action.h>

#include <library/cpp/logger/global/global.h>

#include <util/stream/str.h>
#include <util/generic/string.h>
#include <util/generic/ptr.h>
#include <util/generic/set.h>
#include <saas/indexerproxy/configs/config.h>
#include <saas/protos/tracer.pb.h>

namespace NDispatchableDocument {
    class TBackendTrace {
    private:
        bool IsStarted;
        bool IsFinished;
    public:
        TBackendTrace(const NRTYServer::TBackendTracer& backendTraces);
        TBackendTrace();

        void Serialize(NRTYServer::TBackendTracer* backendTraces) const;

        bool GetIsStarted() const { return IsStarted; }
        bool GetIsFinished() const { return IsFinished; }
        void SetIsStarted(bool Value) { IsStarted = Value; }
        void SetIsFinished(bool Value) { IsFinished = Value; }
    };

    class TTrace {
    private:
        THashMap<TString, TBackendTrace> Backends;
    public:
        void Serialize(NRTYServer::TDocumentTracer* docTracer) const;
        void Deserialize(const NRTYServer::TDocumentTracer* docTracer);

        bool GetTrace(const TString& backend, const TBackendTrace** trace) const;
        TBackendTrace* MutableTrace(const TString& backend);

        bool IsFinished() const;
    };

    struct TMetric {
        ui64 TasksCount = 0;
        ui64 DocsCount = 0;
    };

    NDispatchableDocument::TMetric GetMetric();
};

class TDispatchableDocument {
public:
    typedef TAtomicSharedPtr<TDispatchableDocument> TPtr;
private:
    NSaas::TAction Action;
    ISenderTask& Task;
public:
    TDispatchableDocument(ISenderTask& task);
    ~TDispatchableDocument();

    NRTYServer::TMessage* MutableMessage() {
        return &Action.MutableProtobuf();
    }
    const NRTYServer::TMessage& GetMessage() const {
        return Action.ToProtobuf();
    }

    const NSaas::TAction& GetAction() const {
        return Action;
    }
    NSaas::TAction& GetAction() {
        return Action;
    }

    void PatchMessageByContext(const TProxyServiceConfig& config);
    bool AddInTrace(const NDispatchableDocument::TTrace* startTrace, NRTYServer::TRequestTracer& tracer, const TVector<ISrvDispContext::TPtr>& srvContexts) const;
};
