#pragma once

#include <saas/library/report_builder/abstract.h>
#include <saas/library/searchserver/replier.h>
#include <saas/library/sharding/sharding.h>

#include <search/idl/meta.pb.h>
#include <search/idl/events.ev.pb.h>

#include <library/cpp/object_factory/object_factory.h>
#include <library/cpp/eventlog/eventlog.h>
#include <library/cpp/unistat/unistat.h>

#include <util/string/util.h>
#include <util/string/vector.h>
#include <util/generic/map.h>
#include <util/generic/ptr.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>
#include <util/generic/set.h>
#include <util/generic/hash.h>

class TServiceConfig;

namespace NScatter {
    class ISource;
}

namespace NProxyMeta {
    class TSourceReply {
    public:
        TAtomicSharedPtr<NMetaProtocol::TReport> Report;
        TString Source;

        TSourceReply()
            : Report(new NMetaProtocol::TReport())
        {
        }

        TSourceReply(TAtomicSharedPtr<NMetaProtocol::TReport> report, const TString& source)
            : Report(report)
            , Source(source)
        {
        }
    };

    using TKeysForSource = TMap<NScatter::ISource*, TVector<TString>>;
}

struct TCustomRearrangeParams {
    using TParameters = TMap<TString, TString>;

    TString Name;
    TParameters Parameters;

    bool Init(const TString& name, const TString& configString);
};


class ICustomRearrange {
public:
    virtual ~ICustomRearrange() {}

    bool ProcessSharding(IReplyContext& context,
                         const NSaas::TShardsDispatcher& shardDispatcher,
                         const NSaas::TShardIntervals& shardIntervals,
                         const TVector<THolder<NScatter::ISource>>& sources,
                         NProxyMeta::TKeysForSource& keysForSource,
                         TSelfFlushLogFrame* eventLogFrame) const;

    void ProcessReplies(TVector<NProxyMeta::TSourceReply>& replies,
                        ICustomReportBuilder& builder,
                        IReplyContext& context,
                        TSelfFlushLogFrame* eventLogFrame) const;

    void FormCgiParams(TCgiParameters& /*cgi*/, TSelfFlushLogFrame* /*eventLogFrame*/);

public:
    using TPtr = TAtomicSharedPtr<ICustomRearrange>;


private:
    virtual bool DoFormCgiParameters(TCgiParameters& /*cgi*/) { return false; }

    virtual bool DoProcessSharding(IReplyContext& /*context*/,
                                   const NSaas::TShardsDispatcher& /*shardDispatcher*/,
                                   const NSaas::TShardIntervals& /*shardIntervals*/,
                                   const TVector<THolder<NScatter::ISource>>& /*sources*/,
                                   NProxyMeta::TKeysForSource& /*keysForSource*/) const { return false; }

    virtual void DoProcessReplies(TVector<NProxyMeta::TSourceReply>& /*replies*/,
                                  ICustomReportBuilder& /*builder*/,
                                  IReplyContext& /*context*/) const {}

    virtual TString GetName() const = 0;
};

class ICustomRearrangeFactory {
public:
    using TPtr = THolder<ICustomRearrangeFactory>;
    virtual ~ICustomRearrangeFactory() = default;
    virtual ICustomRearrange::TPtr CreateRerrangeInstance() const = 0;
    virtual TString GetName() const = 0;

    using TFactory = NObjectFactory::TParametrizedObjectFactory<ICustomRearrangeFactory, TString, const TCustomRearrangeParams&>;
};

class TRearrangeEngine {
public:
    TRearrangeEngine(TVector<ICustomRearrange::TPtr>&& rearranges, TSelfFlushLogFrame* eventLogFrame);

    void FormCgiParams(TCgiParameters& cgi);

    bool ProcessSharding(IReplyContext& context,
        const NSaas::TShardsDispatcher& shardDispatcher,
        const NSaas::TShardIntervals& shardIntervals,
        const TVector<THolder<NScatter::ISource>>& sources,
        NProxyMeta::TKeysForSource& keysForSource
    ) const;

    void ProcessReplies(TVector<NProxyMeta::TSourceReply>& replies,
        ICustomReportBuilder& builder, IReplyContext& context) const;

private:
    TVector<ICustomRearrange::TPtr> Rearranges;
    mutable TSelfFlushLogFrame* EventLogFrame;
};

class TRearrangeEngineFactory {
public:
    TRearrangeEngineFactory(const TVector<TCustomRearrangeParams>& params);

    TRearrangeEngine CreateEngine(const TCgiParameters& parameters, TSelfFlushLogFrame* eventLogFrame) const;

private:
    TVector<std::pair<TString, ICustomRearrangeFactory::TPtr>> Factories;
};

/// Interface for custom signals registration in a custom rearrangement implementation
class ICustomRearrangeSignals {
public:
    using TPtr = THolder<ICustomRearrangeSignals>;
    virtual ~ICustomRearrangeSignals() = default;
    virtual void DrillCustomSignals(const TString& service, TUnistat& unistat) const = 0;

    using TFactory = NObjectFactory::TParametrizedObjectFactory<ICustomRearrangeSignals, TString, const TCustomRearrangeParams&>;
};

/// Creates and drills custom rearrangement signals
class TRearrangeSignals {
public:
    void CreateServiceSignals(const TString& service, const TVector<TCustomRearrangeParams>& params);

    void DrillServiceSignals(const TString& service, TUnistat& unistat) const;

private:
    THashMap<TString, TVector<ICustomRearrangeSignals::TPtr>> Instances;
};
