#include "yasm_stats.h"

#include <library/cpp/iterator/cartesian_product.h>

TYasmStats::TYasmStats()
    : WatchedSchemes_({"https"})
    , WatchedHosts_({
        {"www.google.com", "google"},
        {"www.google.ru", "google"},

        {"an.yandex.ru", "yandex"},
        {"www.yandex.by", "yandex"},
        {"www.yandex.com.tr", "yandex"},
        {"www.yandex.com", "yandex"},
        {"www.yandex.kz", "yandex"},
        {"www.yandex.ru", "yandex"},
        {"www.yandex.ua", "yandex"},
        {"www.yandex.uz", "yandex"},
        {"yabs.yandex.ru", "yandex"},
        {"yandex.by", "yandex"},
        {"yandex.com.tr", "yandex"},
        {"yandex.com", "yandex"},
        {"yandex.fr", "yandex"},
        {"yandex.kz", "yandex"},
        {"yandex.ru", "yandex"},
        {"yandex.ua", "yandex"},
        {"yandex.uz", "yandex"},

        {"yastatic.net", "yastatic"},
    })
    , WatchedASNs_({
        "AS12389",
        "AS8359",
        "AS16345",
        "AS8402",
        "AS25513",
        "AS12958",
        "AS3216",
        "AS25159",
        "AS31133",
        "AS6697",
        "AS9198",
        "AS29497",
    })
    , MaxAcceptedAge_(TDuration::Seconds(100))
{
    for (const auto& [scheme, hostTuple, status] : CartesianProduct(
        WatchedSchemes_, WatchedHosts_, TBeacon::AllLegacyStatuses())) {

        Stats_.emplace(MakeSignalName(scheme, hostTuple.first, status), 0);
    }

    THashSet<TString> uniqueHostTypes;
    for (const auto& [host, hostType] : WatchedHosts_) {
        uniqueHostTypes.emplace(hostType);
    }
    for (const auto& [asn, hostType, status] : CartesianProduct(
        WatchedASNs_, uniqueHostTypes, TBeacon::AllLegacyStatuses())) {

        Stats_.emplace(MakeASNSignalName(asn, hostType, status), 0);
        Stats_.emplace(MakeTcpInfoRttSignalName(asn), 0);
    }

    Stats_.emplace(MakeTcpInfoTotalSignalName(), 0);
    Stats_.emplace(MakeTcpInfoTotalPassedSignalName(), 0);
}

void TYasmStats::CountBeacon(const TBeacon& beacon, std::optional<NGeobase::NImpl::StrList> clientASNList) {
    if (beacon.Age > MaxAcceptedAge_)
        return;
    if (!WatchedSchemes_.contains(beacon.Url.GetSchemeInfo().Str))
        return;
    if (!WatchedHosts_.contains(beacon.Url.GetHost()))
        return;
    {
        auto signalName = MakeSignalName(
            beacon.Url.GetSchemeInfo().Str,
            beacon.Url.GetHost(),
            beacon.LegacyStatus
        );
        AtomicAdd(Stats_.at(signalName), beacon.Weight);
    }

    if (!clientASNList)
        return;
    for (const auto& asn : *clientASNList) {
        if (!WatchedASNs_.contains(asn))
            continue;
        auto signalName = MakeASNSignalName(
            asn,
            WatchedHosts_.at(beacon.Url.GetHost()),
            beacon.LegacyStatus
        );
        AtomicAdd(Stats_.at(signalName), beacon.Weight);
    }
}

void TYasmStats::CountTcpInfo(const TTcpInfo& tcpInfo, std::optional<NGeobase::NImpl::StrList> clientASNList) {
    auto totalSignalName = MakeTcpInfoTotalSignalName();
    AtomicAdd(Stats_.at(totalSignalName), 1);

    if (!clientASNList)
        return;

    auto totalPassedSignalName = MakeTcpInfoTotalPassedSignalName();
    AtomicAdd(Stats_.at(totalPassedSignalName), 1);

    for (const auto& asn : *clientASNList) {
        if (!WatchedASNs_.contains(asn))
            continue;
        auto signalName = MakeTcpInfoRttSignalName(asn);
        AtomicAdd(Stats_.at(signalName), tcpInfo.Rtt * 1000);
    }
}

TString TYasmStats::MakeSignalName(TStringBuf scheme, TStringBuf host, TStringBuf status) {
    return TString::Join(scheme, "_", host, "_", status, "_summ");
}

TString TYasmStats::MakeTcpInfoRttSignalName(TStringBuf asn) {
    return TString::Join("rtt_", asn, "_summ");
}

TString TYasmStats::MakeTcpInfoTotalSignalName() {
    return TString::Join("rtt_total_", "received", "_summ");
}

TString TYasmStats::MakeTcpInfoTotalPassedSignalName() {
    return TString::Join("rtt_total_", "passed", "_summ");
}

TString TYasmStats::MakeASNSignalName(TStringBuf asn, TStringBuf hostType, TStringBuf status) {
    // ASNs are in the form of AS0000, so these do not collide with regular signals
    return TString::Join(asn, "_", hostType, "_", status, "_summ");
}
