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

#include <saas/library/sharding/rules/intervals_enumeration.h>

#include <kernel/querydata/saas/key/qd_saas_trie_key.h>

#include <util/string/split.h>

namespace NSaas {
    namespace {
        using NSearchMapParser::TShardIndex;

        const TString& GetUrlMaskPrefix() {
            static const TString urlMaskPrefix = ([] {
                TString urlMaskKey;
                NQueryDataSaaS::TSaaSTrieKey::GenerateMainKey(urlMaskKey, {NQueryDataSaaS::ESaaSSubkeyType::SST_IN_KEY_URL_MASK});
                urlMaskKey.push_back(NQueryDataSaaS::TSaaSTrieKey::Delimiter);
                return urlMaskKey;
            })();
            return urlMaskPrefix;
        }

        TStringBuf GetFirstTwoTokens(TStringBuf s, char delimiter) {
            size_t firstDelim = s.find(delimiter);
            Y_ENSURE(firstDelim != TStringBuf::npos, "key is not a valid QueryDataSaaS key");
            size_t secondDelim = s.find(delimiter, firstDelim + 1);
            if (secondDelim == TStringBuf::npos) {
                secondDelim = s.size();
            }
            return TStringBuf(s.data(), secondDelim);
        }

        TShardIndex GetQuerySearchShardIndex(TStringBuf url, TShardIndex maxIndex) {
            TStringBuf shardKey = GetFirstTwoTokens(url, NQueryDataSaaS::TSaaSTrieKey::Delimiter);
            if (shardKey.StartsWith(GetUrlMaskPrefix())) {
                shardKey = shardKey.Before('/');
            }
            return static_cast<TShardIndex>(CityHash64(shardKey)) % maxIndex;
        }
    }

    struct TQuerySearchRule : TShardsDispatcher::IShardingRule {
        using TShardsDispatcher::IShardingRule::IShardingRule;

        NSearchMapParser::TShardIndex GetShard(const TStringBuf& url, TKeyPrefix /*kps*/) const final {
            return GetQuerySearchShardIndex(url, Context.GetShardsMax());
        }

        bool CheckSearchInterval(TStringBuf url, TKeyPrefix /*kps*/, const TInterval<NSearchMapParser::TShardIndex>& interval) const final {
            const auto shardIndex = GetQuerySearchShardIndex(url, Context.GetShardsMax());
            return CheckInterval(shardIndex, interval);

        }

        bool CheckMessage(const NRTYServer::TMessage& /*message*/, TString& /*error*/) const final {
            return true;
        }

        void EnumerateIntervals(TStringBuf url, TKeyPrefix /*kps*/, const TShardIntervals& sortedIntervals, IShardIntervalCallback& callback) const final {
            const auto shardIndex = GetQuerySearchShardIndex(url, Context.GetShardsMax());
            UrlEnumeration::EnumerateIntervals(shardIndex, sortedIntervals, callback);
        }
    };

    namespace {
        TShardsDispatcher::IShardingRule::TFactory::TRegistrator<TQuerySearchRule> registerThis(QuerySearch);
    }
}
