#pragma once

#include <saas/util/types/shard_index.h>
#include <saas/protos/cluster.pb.h>

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

#include <util/generic/ylimits.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>
#include <library/cpp/cgiparam/cgiparam.h>
#include <library/cpp/json/writer/json_value.h>


namespace NRTYServer {
    class TMessage;
}

namespace NSaas {
    typedef ui64 TKeyPrefix;

    enum ShardingType {
        KeyPrefix /* "keyprefix" */,
        UrlHash /* "url_hash" */,
        External /* "external" */,
        Broadcast /* "broadcast" */,
        Geo /* "geo" */,
        GeoRestrict /* "geo_restrict" */,
        UrlHashErasure /* "url_hash_erasure" */,
        QuerySearch /* "querysearch" */,
        UrlToLastOctothorp /* "url_to_last_octothorp" */,
        UrlToLastUnderscore /* "url_to_last_underscore" */,
        Key = KeyPrefix /* "key" */,
        Url = UrlHash /* "url" */
    };
    enum class ShardsCount {
        Legacy /* "legacy" */,
        UI32 /* "ui32" */
    };

    using TShardIntervals = TVector<NSearchMapParser::TShardsInterval>;

    struct IShardDispatcher {
        virtual ~IShardDispatcher() = default;
        virtual NSearchMapParser::TShardIndex GetShard(const NRTYServer::TMessage& message) const = 0;
        virtual NSearchMapParser::TShardIndex GetShard(const TStringBuf& url, TKeyPrefix kps) const = 0;
    };

    struct IShardIntervalCallback {
        virtual ~IShardIntervalCallback() = default;
        virtual void OnShardInterval(size_t intervalIndex) = 0;
    };

    struct TShardsDispatcher : IShardDispatcher {
    public:
        using TPtr = TAtomicSharedPtr<TShardsDispatcher>;

        struct TContext {
            ShardingType Type = UrlHash;
            ShardsCount Shards = ShardsCount::Legacy;
            ui32 KpsShift = 0;

            TContext(ShardingType type, ShardsCount shards, ui32 kpsShift);
            TContext(ShardingType type, ui32 kpsShift);
            explicit TContext(ShardingType type);

            NSearchMapParser::TShardIndex GetShardsMax() const;

            TString ToString() const;
            void ToProto(NSaasProto::TService&) const;
            static TContext FromString(const TString&);
            static TContext FromProto(const NSaasProto::TService&);

        private:
            TContext() {}
        };

        TShardsDispatcher(const TContext& context);

        NSearchMapParser::TShardIndex GetShard(const NRTYServer::TMessage& message) const override;
        NSearchMapParser::TShardIndex GetShard(const TStringBuf& url, TKeyPrefix kps) const override;
        bool CheckMessage(const NRTYServer::TMessage& message, TString& error) const;
        bool CheckInterval(NSearchMapParser::TShardIndex shard, const TInterval<NSearchMapParser::TShardIndex>& interval) const;
        bool CheckSearchInterval(TStringBuf url, TKeyPrefix kps, const TInterval<NSearchMapParser::TShardIndex>& interval) const;
        bool NeedFullCoverage() const;

        void EnumerateIntervals(TStringBuf url, TKeyPrefix kps, const TShardIntervals& sortedIntervals,
                                const std::function<void (size_t intervalIndex)>& callback) const;
        void EnumerateIntervals(TStringBuf url, TKeyPrefix kps, const TShardIntervals& sortedIntervals,
                                IShardIntervalCallback& interval) const;

        const TContext& GetContext() const;

        class IShardingRule {
        public:
            IShardingRule(const TShardsDispatcher::TContext& context);
            virtual ~IShardingRule() {}

            virtual NSearchMapParser::TShardIndex GetShard(const NRTYServer::TMessage& message) const;
            virtual NSearchMapParser::TShardIndex GetShard(const TStringBuf& url, TKeyPrefix kps) const = 0;
            virtual bool CheckMessage(const NRTYServer::TMessage& message, TString& error) const = 0;
            virtual bool CheckInterval(NSearchMapParser::TShardIndex shard, const TInterval<NSearchMapParser::TShardIndex>& interval) const;
            virtual bool CheckSearchInterval(TStringBuf url, TKeyPrefix kps, const TInterval<NSearchMapParser::TShardIndex>& interval) const = 0;
            virtual bool NeedFullCoverage() const;

            virtual void EnumerateIntervals(TStringBuf url, TKeyPrefix kps, const TShardIntervals& sortedIntervals,
                                            IShardIntervalCallback& callback) const = 0;

            const TContext& GetContext() const;

        public:
            using TFactory = NObjectFactory::TParametrizedObjectFactory<IShardingRule, ShardingType, TContext>;

        protected:
            const TContext Context;
        };

    private:
        THolder<IShardingRule> Impl;
    };

    struct TShardingInfo: public TInterval<NSearchMapParser::TShardIndex> {
    public:
        bool IsValid() const {
            return GetMax() && GetMax() >= GetMin();
        }
    };

    struct TSlotInfo {
        TString ServiceType;
        TString CType;
        TString Service;
        TString Slot;
        TString ConfigType;
        NSearchMapParser::TShardIndex ShardMin;
        NSearchMapParser::TShardIndex ShardMax;
        TString DC;
        bool DisableIndexing = false;
        bool DisableSearch = false;
        bool DisableSearchFiltration = false;
        bool DisableFetch = false;

        bool IsIntSearch = false;

        TString ToString() const;
        NJson::TJsonValue SerializeToJson() const;
        bool DeserializeFromJson(const NJson::TJsonValue& json);
        bool FromString(const TString& string);

        TSlotInfo() {
            ConfigType = "default";
        }
    };

    class TSharding {
    public:
        static TVector<NSearchMapParser::TShardsInterval> SplitInterval(const NSearchMapParser::TShardsInterval& interval, ui32 parts);
        static NSearchMapParser::TShardsInterval GetInterval(ui32 index, ui32 count);
        static NSearchMapParser::TShardsInterval GetInterval(ui32 index, ui32 count, ui32 shardMin, ui32 shardMax);
    };

    NSearchMapParser::TShardIndex GetShards(ShardsCount count);
    TString GetShardString(const NSearchMapParser::TShardIndex shard);
}
