#pragma once

#include "config.h"

#include <saas/searchproxy/common/abstract.h>
#include <saas/searchproxy/configs/serviceconfig.h>
#include <saas/searchproxy/simple_meta/search.h>
#include <saas/searchproxy/proxy_meta/rearrange/abstract/rearrange.h>

#include <saas/library/sharding/sharding.h>
#include <saas/library/searchserver/client.h>

#include <search/meta/scatter/async/runner.h>
#include <search/meta/scatter/connstat.h>
#include <search/meta/scatter/result.h>
#include <search/meta/scatter/task.h>
#include <search/meta/scatter/source.h>

#include <infra/yp_service_discovery/libs/sdlib/client.h>
#include <library/cpp/mediator/messenger.h>

namespace NProxyMeta {
    class TSearch: public NSimpleMeta::TBaseProxySearch {
    private:
        NSaas::TShardIntervals ShardIntervals;
        TVector<THolder<NScatter::ISource>> Sources;
        TRearrangeEngineFactory RearrangeFactory;
        CPHashFunction HashFunction = nullptr;
        THolder<NScatter::IConnStatMap> ConnStatMap;
        THolder<NYP::NServiceDiscovery::TSDClient> SDClient;

    public:
        using TPtr = TAtomicSharedPtr<TSearch>;

        static TString GetClassName() {
            return "proxy";
        }

    public:
        TSearch(const TServiceConfig& serviceConfig, const TSearchSources& searchSources
            , const NSaas::TShardsDispatcher::TPtr sharding, ICgiCorrectorPtr cgiCorrector
            , const TServiceDiscoveryOptions& sdOpts, const TString& serviceName
            , TEventLogPtr eventlog);

        virtual bool DoSearchContext(IReplyContext::TPtr context) const override;

        const NProxyMeta::TConfig& GetConfig() const {
            CHECK_WITH_LOG(ServiceConfig.GetProxyMetaConfig());
            return *ServiceConfig.GetProxyMetaConfig();
        }

    private:
        void DefaultSharding(IReplyContext& context, bool noTextSplit, TKeysForSource& keysForSource) const;

    };

    class TProxy: public IProxy, public IMessageProcessor {
    private:
        TMap<TString, TSearch::TPtr> Searchers;
        TMap<TString, TString> SearchersNotStarted;

        static TFactory::TRegistrator<TProxy> Registrator;

    public:
        virtual ISearchReplier::TPtr BuildReplier(const TString& service, IReplyContext::TPtr context) override;

        virtual void Start() override {
            for (auto&& i : Searchers) {
                try {
                    i.second->Start();
                } catch (const yexception& e) {
                    SearchersNotStarted[i.first] = e.what();
                    ERROR_LOG << "Proxy_meta failed for " << i.first << ": " << e.what() << Endl;
                }
            }
            RegisterGlobalMessageProcessor(this);
        }

        virtual void Stop() override {
            UnregisterGlobalMessageProcessor(this);
            for (auto&& i : Searchers) {
                i.second->Stop();
            }
        }

        virtual void ReopenLogs() override {
        }

        virtual TMap<TString, TString>& GetFailedSearchers() override {
            return SearchersNotStarted;
        }

        virtual void RegisterService(const NSearchMapParser::TSearchMapService& serviceInfo, const TServiceConfig& serviceConfig,
            const TSourcesConfig& sourcesConfig, ICgiCorrectorPtr cc, TServiceMetricPtr metric, TEventLogPtr eventlog) override {

            Y_UNUSED(metric);
            if (serviceConfig.GetProxyMetaConfig()) {
                Searchers[serviceInfo.Name] = new TSearch(serviceConfig, sourcesConfig.SearchSources, serviceInfo.ShardsDispatcher,
                    cc, sourcesConfig.ServiceDiscoveryOptions, serviceInfo.Name, eventlog);
                INFO_LOG << "Registered ProxyMeta for " << serviceInfo.Name << Endl;
            }
        }

        virtual void RegisterMetaService(const NSearchMapParser::TMetaService& serviceInfo, const TServiceConfig& serviceConfig, const TSourcesConfig& sourcesConfig
                , ICgiCorrectorPtr cc, TServiceMetricPtr metric, TEventLogPtr eventlog) override {
            Y_UNUSED(metric);
            if (serviceConfig.GetProxyMetaConfig()) {
                NSaas::TShardsDispatcher::TPtr sharding = MakeAtomicShared<NSaas::TShardsDispatcher>(
                    NSaas::TShardsDispatcher::TContext(NSaas::Broadcast)
                );
                Searchers[serviceInfo.Name] = new TSearch(serviceConfig, sourcesConfig.SearchSources, sharding, cc, TServiceDiscoveryOptions(), serviceInfo.Name, eventlog);
                INFO_LOG << "Registered ProxyMeta for " << serviceInfo.Name << Endl;
            }
        }

        virtual void RegisterUpperService(const TServiceConfig& /*serviceConfig*/, const TSourcesConfig& /*sourcesConfig*/, ICgiCorrectorPtr /*cc*/
                , TServiceMetricPtr /*metric*/, TEventLogPtr /*eventlog*/) override {
            return;
        }

        virtual bool Process(IMessage* message) override final;

        virtual TString Name() const override final {
            return "proxy_meta";
        }
    };
}
