#include "erasure.h"

#include <saas/protos/rtyserver.pb.h>
#include <library/cpp/logger/global/global.h>

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

namespace NSaas {

    namespace {
        constexpr TStringBuf erasurePrefix("erasure_", sizeof("erasure_") - 1);
        constexpr ui32 numberWidth = 2;
        constexpr ui32 plusPosition = erasurePrefix.size() + numberWidth;
        constexpr ui32 secondUnderlinePosition = plusPosition + numberWidth + 1;
        constexpr ui32 thirdUnderlinePosition = secondUnderlinePosition + numberWidth + 1;
        constexpr ui32 firstNumberPosition = erasurePrefix.size();
        constexpr ui32 secondNumberPosition = firstNumberPosition + numberWidth + 1;
        constexpr ui32 thirdNumberPosition = secondNumberPosition + numberWidth + 1;
        constexpr ui32 totalPrefixSize = erasurePrefix.size() + (numberWidth + 1) * 3;
    }

    // ========= TErasureUrlConverter =========
    bool TErasureUrlConverter::InitializeFromBaseUrl(TStringBuf url) {
        if (url.empty()) {
            return false;
        }
        baseUrl = url;
        return true;
    }

    TStringBuf TErasureUrlConverter::GetBaseUrl() const {
        return baseUrl;
    }

    ui32 TErasureUrlConverter::GetMainNumber() const {
        return mainNumber;
    }

    ui32 TErasureUrlConverter::GetParityNumber() const {
        return parityNumber;
    }

    ui32 TErasureUrlConverter::GetPartNumber() const {
        return partNumber;
    }

    bool TErasureUrlConverter::SetErasureParams(ui32 main, ui32 parity) {
        if (main >= 100 || main <= 0 || parity <=0 || parity >= 100) {
            return false;
        }

        mainNumber = main;
        parityNumber = parity;
        return true;
    }

    bool TErasureUrlConverter::InitializeFromFormattedUrl(TStringBuf url) {
        // Format "erasure_xx+yy_zz_*"
        // Checking size
        if (url.size() <= totalPrefixSize) {
            return false;
        }
        bool result = true;
        result = result && url.substr(0, erasurePrefix.size()) == erasurePrefix;
        result = result && url[plusPosition] == '+';
        result = result && url[secondUnderlinePosition] == '_';
        result = result && url[thirdUnderlinePosition] == '_';
        // Extracting numbers
        result = result && TryIntFromString<10, ui32>(url.substr(firstNumberPosition, numberWidth), mainNumber);
        result = result && TryIntFromString<10, ui32>(url.substr(secondNumberPosition, numberWidth), parityNumber);
        result = result && TryIntFromString<10, ui32>(url.substr(thirdNumberPosition, numberWidth), partNumber);
        // Extracting base url
        baseUrl = url.substr(totalPrefixSize, url.size() - totalPrefixSize);
        return result;
    }

    TString TErasureUrlConverter::ConvertNumber(ui32 number) const {
        Y_ENSURE(number > 0 && number < 100);
        if (number < 10) {
            return "0" + IntToString<10>(number);
        } else {
            return IntToString<10>(number);
        }
    }

    TString TErasureUrlConverter::GetFormattedUrl(ui32 part) {
        if (part < 0 || part >= 100 || part >= mainNumber + parityNumber) {
            return TString("");
        }

        return TString::Join("erasure_", ConvertNumber(mainNumber), "+", ConvertNumber(parityNumber), "_", ConvertNumber(part), "_", TString{baseUrl});
    }

    // ========= TUrlErasureShardingRule ==========
    bool TUrlErasureShardingRule::CheckMessage(const NRTYServer::TMessage& message, TString& error) const {
        if (!message.GetDocument().HasUrl()) {
            error = "Document doesn't have url";
            return false;
        }
        return true;
    }

    bool TUrlErasureShardingRule::ParseUrl(TStringBuf url, TStringBuf& baseUrl, ui32& mainNumber, ui32& parityNumber, ui32& partNumber) const {
        TErasureUrlConverter converter;
        auto result = converter.InitializeFromFormattedUrl(url);
        if (result) {
            baseUrl = converter.GetBaseUrl();
            mainNumber = converter.GetMainNumber();
            parityNumber = converter.GetParityNumber();
            partNumber = converter.GetPartNumber();
        }
        return result;
    }

    NSearchMapParser::TShardIndex TUrlErasureShardingRule::GetUrlErasureShard(const TStringBuf& url, const NSearchMapParser::TShardIndex shards) const {
        TStringBuf baseUrl;
        ui32 mainNumber;
        ui32 parityNumber;
        ui32 partNumber;
        if (!ParseUrl(url, baseUrl, mainNumber, parityNumber, partNumber)) {
            Y_ENSURE(false);
        }

        ui64 hash = CityHash64(baseUrl);
        ui32 totalNumber = mainNumber + parityNumber;
        auto roundShards = shards - shards % totalNumber;
        auto step = roundShards / totalNumber;

        Y_ENSURE(totalNumber >= partNumber);

        NSearchMapParser::TShardIndex shardIndex = (static_cast<NSearchMapParser::TShardIndex>(hash) + step * partNumber) % roundShards;
        return shardIndex;
    }

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

    bool TUrlErasureShardingRule::CheckSearchInterval(TStringBuf url, TKeyPrefix /*kps*/, const TInterval<NSearchMapParser::TShardIndex>& interval) const {
        if (url.empty()) {
            return true;
        }

        const NSearchMapParser::TShardIndex shard = GetUrlErasureShard(url, Context.GetShardsMax());
        return CheckInterval(shard, interval);
    }

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

    TShardsDispatcher::IShardingRule::TFactory::TRegistrator<TUrlErasureShardingRule> TUrlErasureShardingRule::Registrator(UrlHashErasure);
}
