#pragma once

#include "urlhash.h"

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

namespace NSaas {
    template <char separator>
    class TUrlToLastCharShardingRule : public TShardsDispatcher::IShardingRule {
        using TSelf = TUrlToLastCharShardingRule<separator>;

    public:
        TUrlToLastCharShardingRule(const TShardsDispatcher::TContext& context)
            : TShardsDispatcher::IShardingRule(context)
            , BaseRule(context)
        {
        }

        virtual bool CheckMessage(const NRTYServer::TMessage& message, TString& error) const override {
            return BaseRule.CheckMessage(message, error);
        }

        virtual NSearchMapParser::TShardIndex GetShard(const TStringBuf& url, TKeyPrefix kps) const override {
            return BaseRule.GetShard(GetUrlPrefix(url), kps);
        }

        virtual bool CheckSearchInterval(TStringBuf url, TKeyPrefix kps, const TInterval<NSearchMapParser::TShardIndex>& interval) const override {
            return BaseRule.CheckSearchInterval(GetUrlPrefix(url), kps, interval);
        }

        virtual void EnumerateIntervals(TStringBuf url, TKeyPrefix kps, const TShardIntervals& sortedIntervals,
                                        IShardIntervalCallback& callback) const override {
            return BaseRule.EnumerateIntervals(GetUrlPrefix(url), kps, sortedIntervals, callback);
        }

    private:
        static TStringBuf GetUrlPrefix(TStringBuf url) {
            TStringBuf prefix, suffix;
            if (url.TryRSplit(separator, prefix, suffix)) {
                return prefix;
            }
            return url;
        }

    private:
        TUrlShardingRule BaseRule;

        static TShardsDispatcher::IShardingRule::TFactory::TRegistrator<TSelf> Registrator;
    };

    // using the world "octothorp" instead of "hash" to avoid confusion with hash functions
    using TUrlToLastOctothorpShardingRule = TUrlToLastCharShardingRule<'#'>;

    using TUrlToLastUnderscoreShardingRule = TUrlToLastCharShardingRule<'_'>;
}

