#include "sharding.h"

#include <library/cpp/logger/global/global.h>

#include <util/generic/strbuf.h>
#include <util/string/cast.h>
#include <util/string/split.h>
#include <util/system/yassert.h>

#include <cctype>

NRTYServer::TShard NRTYServer::GetShard(ui64 keyprefix, NRTYServer::TShard shardsNumber) {
    return (keyprefix && shardsNumber) ? static_cast<TShard>(keyprefix) % shardsNumber : 0;
}

NRTYServer::TShard NRTYServer::GetShard(const TString& directory) {
    TString dirName(TFsPath(directory).Basename());
    auto numberPos = dirName.find('0');
    Y_ASSERT(numberPos != TString::npos);
    auto numberBegin = dirName.begin() + numberPos;
    auto numberEnd = numberBegin;
    while (numberEnd != dirName.end() && std::isdigit(*numberEnd)) {
        ++numberEnd;
    }
    NRTYServer::TShard id = 0;
    TryFromString(TStringBuf(numberBegin, numberEnd), id);
    return id;
}

bool NRTYServer::GetShardAndSegmentId(const TString& directory, size_t& shard, i32& segmentId) {
    const TString dirName(TFsPath(directory).Basename());
    size_t offset = 0;
    if (dirName.StartsWith("index_")) {
        offset = strlen("index_");
    } else if (dirName.StartsWith("memory_index_")) {
        offset = strlen("memory_index_");
    } else if (dirName.StartsWith("temp__")) {
        offset = strlen("temp__");
    } else if (dirName.StartsWith("prep_index_")) {
        offset = strlen("prep_index_");
    }
    if (offset) {
        TVector<TString> parts;
        Split(dirName.data() + offset, "_", parts);
        if (parts.size() == 2) {
            try {
                shard = FromString<size_t>(parts[0]);
                segmentId = FromString<i32>(parts[1]);
                return true;
            } catch (...) {
                ERROR_LOG << CurrentExceptionMessage() << Endl;
            }
        } else {
            return false;
        }
    }
    return false;
}
