#include "pessimization.h"

namespace NSrvKernel::NDynamicBalancing {
    bool IsBackendHealthy(TDynamicBackend& backend) noexcept {
        return backend.Enabled() && backend.WeightFromPing() > 0;
    }
}

using namespace NSrvKernel::NDynamicBalancing;

namespace {
    double Clamp(double value, double min, double max) {
        return Max(Min(value, max), min);
    }

    bool IsBackendIgnored(TDynamicBackendHolder& backend) {
        return backend->Status == TDynamicBackend::TStatus::Blacklisted;
    }

    bool IsBackendDegraded(TDynamicBackend& backend) {
        return backend.Degraded();
    }

    static constexpr double MinWeight = 0;
    static constexpr double MaxWeight = 10;

    using TSupersedingStatusesTable = std::array<
        TVector<TDynamicBackend::TStatus>, ToUnderlying(TDynamicBackend::TStatus::Max) + 1>;

    TSupersedingStatusesTable BuildSupersedingStatusesTable() {
        TSupersedingStatusesTable table;

        table[ToUnderlying(TDynamicBackend::TStatus::NotReady)] = {
            TDynamicBackend::TStatus::NotChecked,
            TDynamicBackend::TStatus::Recovering,
            TDynamicBackend::TStatus::Degraded,
            TDynamicBackend::TStatus::Critical,
            TDynamicBackend::TStatus::Unhealthy,
        };

        table[ToUnderlying(TDynamicBackend::TStatus::Unhealthy)] = {
            TDynamicBackend::TStatus::NotChecked,
            TDynamicBackend::TStatus::Recovering,
            TDynamicBackend::TStatus::Degraded,
            TDynamicBackend::TStatus::Critical,
        };

        table[ToUnderlying(TDynamicBackend::TStatus::Critical)] = {
            TDynamicBackend::TStatus::NotChecked,
            TDynamicBackend::TStatus::Recovering,
            TDynamicBackend::TStatus::Degraded,
        };

        table[ToUnderlying(TDynamicBackend::TStatus::Degraded)] = {
            TDynamicBackend::TStatus::NotChecked,
            TDynamicBackend::TStatus::Recovering,
        };

        return table;
    }

    static auto SupersedingStatuses = BuildSupersedingStatusesTable();

    bool CanApplyCheck(ui64 requests, ui64 threshold, bool& hasSkippedCheck) noexcept {
        if (requests >= threshold && requests > 1) {
            return true;
        }
        hasSkippedCheck = true;
        return false;
    }
}

TDynamicBackend::TStatus TPessimizationTracker::RunBackendChecks(TDynamicBackend& backend, TStateAggregate totalStats, bool& hasFiredCheck, bool& hasSkippedCheck) noexcept {
    hasSkippedCheck = false;
    hasFiredCheck = false;
    if (!IsBackendReady(backend)) {
        return TDynamicBackend::TStatus::NotReady;
    }
    // Failed health check transitions a backend to the unhealthy state.
    if (!IsBackendHealthy(backend)) {
        return TDynamicBackend::TStatus::Unhealthy;
    }

    const auto stats = backend.CurrentStats();
    ui64 requests = stats.Requests;
    ui64 fails = stats.Fails;
    ui64 timeouts = stats.Timeouts;
    ui64 sumProcessingTime = stats.SumProcessingTime;

    bool anyCheckLaunched = false;
    if (Config_.has_fail_rate_check() && CanApplyCheck(requests, Config_.fail_rate_check().min_requests(), hasSkippedCheck)) {
        anyCheckLaunched = true;
        if (1.f * (fails + timeouts) / requests >= Config_.fail_rate_check().threshold()) {
            Log_ << "backend [" << backend.Name() << "]: fail rate check failed" << Endl;
            hasFiredCheck = true;
            return TDynamicBackend::TStatus::Critical;
        }
    }

    if (Config_.has_consecutive_fails_check()) {
        ui64 tracked = backend.TrackedConsecutiveRequests.load(std::memory_order_relaxed);
        ui64 max = backend.ConsecutiveFailsMax.load(std::memory_order_relaxed);
        if (max >= Config_.consecutive_fails_check().threshold()) {
            Log_ << "backend [" << backend.Name() << "]: consecutive fails check failed" << Endl;
            backend.ResetConsecutiveFails();
            hasFiredCheck = true;
            return TDynamicBackend::TStatus::Critical;
        }
        if (tracked >= Config_.consecutive_fails_check().threshold()) {
            anyCheckLaunched = true;
        } else {
            hasSkippedCheck = true;
        }
    }

    if (Config_.has_response_time_check() && CanApplyCheck(requests, Config_.response_time_check().min_requests(), hasSkippedCheck)) {
        anyCheckLaunched = true;
        double responseTimeAverage = 1. * sumProcessingTime / requests;
        if (Abs(responseTimeAverage - totalStats.TotalResponseTimeAverage) >= Config_.response_time_check().max_deviation().MilliSeconds()) {
            Log_ << "backend [" << backend.Name() << "]: response time check failed" << Endl;
            hasFiredCheck = true;
            return TDynamicBackend::TStatus::Degraded;
        }
    }

    if (IsBackendDegraded(backend)) {
        return TDynamicBackend::TStatus::Degraded;
    }

    if (!anyCheckLaunched) {
        if (backend.Status == TDynamicBackend::TStatus::Unhealthy ||
            backend.Status == TDynamicBackend::TStatus::NotReady ||
            (backend.Status == TDynamicBackend::TStatus::NotChecked && backend.ActiveChecked)) {
            return TDynamicBackend::TStatus::Healthy;
        }
        return backend.Status;
    }

    return TDynamicBackend::TStatus::Healthy;
}

void TPessimizationTracker::LogInitialInfo() const noexcept {
    Log_ << "max pessimized backends = " << MaxPessimizedBackendsCount_ << " (total " << Backends_.size() << ")" << Endl;
}

void TPessimizationTracker::SetBackendStatusAndLog(TDynamicBackendHolder& backend, TDynamicBackend::TStatus newStatus) noexcept {
    TDynamicBackend::TStatus oldStatus = backend->Status;
    if (oldStatus != newStatus) {
        Log_ << "backend [" << backend->Name() << "]: status change " << oldStatus << " -> " << newStatus << Endl;
    }

    backend->Status = newStatus;
}

void TPessimizationTracker::DiscoverNewBackends(const TVector<TBackendDescriptor::TRef>& newBackends, const TPessimizationTracker* donor) {
    TDuration historyInterval = Config_.history_interval();
    ui64 minRequests = 0;
    if (Config_.has_fail_rate_check()) {
        minRequests = Config_.fail_rate_check().min_requests();
    }
    if (Config_.has_response_time_check() && minRequests < Config_.response_time_check().min_requests()) {
        minRequests = Config_.response_time_check().min_requests();
    }

    for (const auto& newBackend : newBackends) {
        TIntrusivePtr<TDynamicBackend> dynamicBackend;
        if (donor) {
            auto it = donor->Backends_.find(newBackend->Name());
            if (it != donor->Backends_.end()) {
                dynamicBackend = it->second->Clone();
            }
        }
        auto holder = MakeHolder<TDynamicBackendHolder>(dynamicBackend ?: MakeIntrusive<TDynamicBackend>(newBackend, minRequests, historyInterval));
        auto* backend = holder.Get();
        backend->Get()->SetWeightNormalizationCoeff(Config_.has_active() ? Config_.active().weight_normalization_coeff() : 1);

        if (!dynamicBackend) {
            // By default, backend is enabled, is not pessimized and its weight set to 1...
            backend->Get()->SetEnabled(true);
            backend->Get()->SetWeightFromPing(1, false);
            backend->Get()->SetDegraded(false);
            backend->Get()->WeightCoeff = 1;
            backend->Get()->TargetWeightCoeff = 1;
            // unless this is SD update, and we have active checks configured...
            if (donor && Config_.has_active()) {
                // then we try to prevent use of new backend prior to first active check
                backend->Get()->Status = TDynamicBackend::TStatus::NotChecked;
            }
        }

        Backends_.emplace(backend->Get()->Name(), std::move(holder));
    }

    MaxPessimizedBackendsCount_ = Backends_.size() * Config_.max_pessimized_share();

    // Try to pessimize again pessimized backends taken from the donor
    // taking into account actual MaxPessimizedBackendsCount_
    if (donor) {
        for (const auto& entry : Backends_) {
            auto& backend = *entry.second.Get();
            if (backend->Pessimized) {
                backend->Pessimized = false;
                backend.Unlink();
                if (!TryPessimizeBackend(backend, false)) {
                    backend->WeightCoeff = 1;
                }
            }
        }
    }
}

void TPessimizationTracker::UpdateNotReadySet(TIntrusivePtr<TNotReadyBackendSet> notReadySet) noexcept {
    NotReadySet_ = std::move(notReadySet);
}

void TPessimizationTracker::UpdateBlacklist(TIntrusivePtr<TBlacklist> blacklist) noexcept {
    Blacklist_ = std::move(blacklist);
}

bool TPessimizationTracker::TrySupersedBackendFor(TDynamicBackend::TStatus newStatus) noexcept {
    for (auto status : SupersedingStatuses[static_cast<size_t>(newStatus)]) {
        auto& pessimizedBackendsForStatus = PessimizedBackends_[static_cast<size_t>(status)];
        if (!pessimizedBackendsForStatus.Empty()) {
            auto& backend = *pessimizedBackendsForStatus.Front();
            Y_ASSERT(backend->Status == status);
            Log_ << "backend [" << backend->Name() << "]: superseded by a new " << newStatus << " backend" << Endl;
            backend->WeightCoeff = 1;
            UnpessimizeBackend(backend);
            return true;
        }
    }
    return false;
}

bool TPessimizationTracker::TryPessimizeBackend(TDynamicBackendHolder& backend, bool needReset) noexcept {
    if (backend->Pessimized) {
        // This case occurs when backend immediately transits from RECOVERING status to the non-HEALTHY status
        // or from any pessimized status to the NotReady or Unhealthy status.
        // Such a backend is already accounted as pessimized and occupies a pessimization slot.
        // Just remove it from PessimizedBackends_. To reinsert later to the appropriate list.
        backend.Unlink();
    } else {
        // Otherwise, we need to allocate a pessimization slot.
        if (PessimizedBackendsCount_ == MaxPessimizedBackendsCount_) {
            // Pessimization limit is exceeded. We either don't pessimize this backend or have to unpessimize some other first.
            OverPessimized_ = true;
            if (!TrySupersedBackendFor(backend->Status)) {
                // Superseding process have failed. We cannot pessimize this backend now.
                return false;
            }
        }
        PessimizedBackendsCount_++;
    }

    backend->Pessimized = true;
    if (needReset) {
        backend->PessimizationStepsRemaining = backend->CurrentPessimizationInterval.load();
        backend->WeightCoeff = 0;
        backend->TargetWeightCoeff = 0;
    }
    PessimizedBackends_[static_cast<size_t>(backend->Status.load())].PushBack(&backend);
    return true;
}

void TPessimizationTracker::RecoverBackend(TDynamicBackendHolder& backend) noexcept {
    Y_ASSERT(backend->Pessimized);
    Y_ASSERT(PessimizedBackendsCount_ > 0);

    // After backend pessimization, it is returned to the balancing state, but counts as pessimized,
    // until weight coefficient becomes greater than TargetWeightCoeff.

    backend->WeightCoeff = 0;
    backend->PessimizationStepsRemaining = 0;
    backend->Pessimized = true;
    SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Recovering);

    backend.Unlink();
    PessimizedBackends_[static_cast<size_t>(backend->Status.load())].PushBack(&backend);

    // After backend returned to the balancing state, its weight should be greater than zero.
    UpdateWeightCoeff(backend);
    backend->DropOldStats();
}

void TPessimizationTracker::UnpessimizeBackend(TDynamicBackendHolder& backend) noexcept {
    Y_ASSERT(backend->Pessimized);
    Y_ASSERT(PessimizedBackendsCount_ > 0);

    backend.Unlink();
    PessimizedBackendsCount_--;
    backend->Pessimized = false;
}

void TPessimizationTracker::UpdateWeightCoeff(TDynamicBackendHolder& backend, bool canIncrease) noexcept {
    Y_ASSERT(!backend->Pessimized || backend->Status == TDynamicBackend::TStatus::Recovering);

    if (Config_.has_active() && Config_.active().use_backend_weight()) {
        double targetWeightCoeff = Clamp(backend->WeightFromPing(), MinWeight, MaxWeight);
        if (targetWeightCoeff > 0) {
            backend->TargetWeightCoeff = targetWeightCoeff;
        }
    } else {
        backend->TargetWeightCoeff = 1;
    }

    if (backend->WeightCoeff < backend->TargetWeightCoeff && canIncrease) {
        double oldWeightCoeff = backend->WeightCoeff;
        backend->WeightCoeff = Min<double>(backend->WeightCoeff + Config_.weight_increase_step(), backend->TargetWeightCoeff);
        Log_ << "backend [" << backend->Name() << "]: weight coeff change "
            << oldWeightCoeff << " -> " << backend->WeightCoeff.load()
            << " (target = " << backend->TargetWeightCoeff.load() << ")" << Endl;
    } else if (backend->WeightCoeff > backend->TargetWeightCoeff) {
        Log_ << "backend [" << backend->Name() << "]: weight coeff change "
            << backend->WeightCoeff.load() << " -> " << backend->TargetWeightCoeff.load() << Endl;
        backend->WeightCoeff = backend->TargetWeightCoeff.load();
    }

    if (backend->Status == TDynamicBackend::TStatus::Recovering && backend->WeightCoeff == backend->TargetWeightCoeff)  {
        if (backend->Pessimized) {
            UnpessimizeBackend(backend);
        }
        SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Healthy);
    }
}

void TPessimizationTracker::UpdateBackend(TDynamicBackendHolder& backend, const TStateAggregate& aggregate) noexcept {
    if (IsBackendIgnored(backend)) {
        Y_ASSERT(false);
        return;
    }

    TDynamicBackend::TStatus oldStatus = backend->Status;

    bool hasSkippedCheck = false;
    bool hasFiredCheck = false;
    auto newStatus = RunBackendChecks(*backend.Get(), aggregate, hasFiredCheck, hasSkippedCheck);

    if (oldStatus == TDynamicBackend::TStatus::Recovering && newStatus == TDynamicBackend::TStatus::Healthy) {
        newStatus = TDynamicBackend::TStatus::Recovering;
    }

    SetBackendStatusAndLog(backend, newStatus);

    switch (newStatus) {
        case TDynamicBackend::TStatus::Healthy:
            // HEALTHY statuses always decreases pessimization interval, if all checks were applied.
            if (!hasSkippedCheck) {
                backend->CurrentPessimizationInterval = Max(3 * backend->CurrentPessimizationInterval / 4, 1);
            }
            break;

        case TDynamicBackend::TStatus::Recovering:
            // RECOVERING status means that backend is not yet good, so we delay CurrentPessimizationInterval decrease until it becomes HEALTHY.
            break;

        case TDynamicBackend::TStatus::Unhealthy:
        case TDynamicBackend::TStatus::NotReady:
            if (TryPessimizeBackend(backend)) {
                Log_ << "backend [" << backend->Name() << "]: pessimized until it comes back" << Endl;
            }
            break;

        case TDynamicBackend::TStatus::Critical:
        case TDynamicBackend::TStatus::Degraded:
            if (TryPessimizeBackend(backend)) {
                Log_ << "backend [" << backend->Name() << "]: pessimized for " << backend->PessimizationStepsRemaining.load() << " steps (" << backend->Status.load() << ")" << Endl;
            }
            // Even if we failed to pessimize backend, increase its pessimization interval, if we applied some check.
            if (hasFiredCheck) {
                backend->CurrentPessimizationInterval = Min<int>(4 * backend->CurrentPessimizationInterval / 3 + 1,
                                                                 Config_.max_pessimization_steps());
            }
            break;

        case TDynamicBackend::TStatus::NotChecked:
            if (TryPessimizeBackend(backend)) {
                Log_ << "backend [" << backend->Name() << "]: pessimized until first active check" << Endl;
            }
            break;

        default:
            Y_VERIFY(false);
    }

    if (!backend->Pessimized || backend->Status == TDynamicBackend::TStatus::Recovering) {
        UpdateWeightCoeff(backend, !hasSkippedCheck);
    }
}

bool TPessimizationTracker::PessimizeIfNotReady(TDynamicBackendHolder& backend) noexcept {
    if (!IsBackendReady(*backend.Get())) {
        SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::NotReady);
        return TryPessimizeBackend(backend);
    }
    return false;
}

bool TPessimizationTracker::PessimizeIfNotHealthy(TDynamicBackendHolder& backend) noexcept {
    if (!IsBackendHealthy(*backend.Get())) {
        SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Unhealthy);
        return TryPessimizeBackend(backend);
    }
    return false;
}

bool TPessimizationTracker::PessimizeIfDegraded(TDynamicBackendHolder& backend) noexcept {
    if (IsBackendDegraded(*backend.Get())) {
        SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Degraded);
        return TryPessimizeBackend(backend);
    }
    return false;
}

void TPessimizationTracker::UpdatePessimizedBackends(TStateAggregate& stateAggregate) noexcept {
    for (const auto& entry : Backends_) {
        auto& backend = *entry.second.Get();
        if (!backend->Pessimized) {
            continue;
        }

        // Ignored backend must not occupy the pessimization slot
        Y_ASSERT(!IsBackendIgnored(backend));

        switch (backend->Status) {
            case TDynamicBackend::TStatus::Recovering:
                Y_ASSERT(backend->PessimizationStepsRemaining == 0);
                UpdateBackend(backend, stateAggregate);
                break;

            case TDynamicBackend::TStatus::Unhealthy:
                if (!PessimizeIfNotReady(backend) && IsBackendHealthy(*backend.Get()) && !PessimizeIfDegraded(backend)) {
                    // UNHEALTHY backend is recovered as fast as possible.
                    RecoverBackend(backend);
                }
                break;

            case TDynamicBackend::TStatus::NotReady:
                if (IsBackendReady(*backend.Get())) {
                    if (!PessimizeIfNotHealthy(backend) && !PessimizeIfDegraded(backend)) {
                        // NotReady backend is recovered as fast as possible.
                        RecoverBackend(backend);
                    }
                }
                break;

            case TDynamicBackend::TStatus::NotChecked:
                if (!PessimizeIfNotReady(backend) && !PessimizeIfNotHealthy(backend) && !PessimizeIfDegraded(backend) && backend->ActiveChecked) {
                    backend->WeightCoeff = 1;
                    UnpessimizeBackend(backend);
                    backend->DropOldStats();
                    SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Healthy);
                }
                break;

            default:
                if (!PessimizeIfNotReady(backend) && !PessimizeIfNotHealthy(backend)) {
                    backend->PessimizationStepsRemaining--;
                    if (backend->PessimizationStepsRemaining == 0) {
                        // RECOVERING backend is accounted as pessimized, so tracker will not update it on this step.
                        // So if weight_increase_step < 1, RECOVERING state will last at least one step.
                        if (!PessimizeIfDegraded(backend)) {
                            RecoverBackend(backend);
                        }
                    }
                }
                break;
        }
    }
}

void TPessimizationTracker::UpdateBlacklistedBackends() noexcept {
    for (const auto& entry : Backends_) {
        auto& backend = *entry.second.Get();
        bool blacklisted = Blacklist_ && Blacklist_->Contains(backend->Name());
        if (blacklisted == (backend->Status == TDynamicBackend::TStatus::Blacklisted)) {
            continue;
        }
        if (backend->Pessimized) {
            backend->WeightCoeff = 1;
            UnpessimizeBackend(backend);
        }
        if (blacklisted) {
            SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Blacklisted);
        } else {
            SetBackendStatusAndLog(backend, TDynamicBackend::TStatus::Healthy);
            backend->DropOldStats();
        }
    }
}

bool TPessimizationTracker::IsBackendReady(const TDynamicBackend& backend) noexcept {
    return !NotReadySet_ || !NotReadySet_->contains(backend.Name());
}

TVector<TIntrusivePtr<TDynamicBackend>> TPessimizationTracker::Step() noexcept {
    Log_.SetName(TStringBuilder{} << Config_.backends_name() << "/s" << Step_);
    Step_++;

    OverPessimized_ = false;

    auto stateAggregate = CalculateStateAggregate();

    UpdateBlacklistedBackends();
    // First of all, update all pessimized backends, because they can instantly transit to the unpessimized step.
    UpdatePessimizedBackends(stateAggregate);

    for (const auto& entry : Backends_) {
        auto& backend = *entry.second.Get();

        if (backend->Pessimized || IsBackendIgnored(backend)) {
            continue;
        }

        UpdateBackend(backend, stateAggregate);
    }

    TVector<TIntrusivePtr<TDynamicBackend>> backendsRef(Reserve(Backends_.size()));
    for (const auto& entry : Backends_) {
        auto& backend = *entry.second.Get();

        if (IsBackendIgnored(backend)) {
            continue;
        }

        // RECOVERING backends are accounted in PessimizedBackendsCount_ (thus they are pessimized), but they are ready to receive traffic.
        if (backend->Pessimized && backend->Status != TDynamicBackend::TStatus::Recovering) {
            continue;
        }

        backendsRef.push_back(backend.Clone());
    }

    return backendsRef;
}

TStateAggregate TPessimizationTracker::CalculateStateAggregate() noexcept {
    TStateAggregate result;
    result.TotalRequests = 0;
    result.TotalResponseTimeAverage = 0;

    if (!StateAggregationsNeeded()) {
        return result;
    }

    for (const auto& entry : Backends_) {
        auto stats = entry.second->Get()->CurrentStats();
        result.TotalRequests += stats.Requests;
        result.TotalResponseTimeAverage += stats.SumProcessingTime;
    }

    if (result.TotalRequests > 0) {
        result.TotalResponseTimeAverage /= result.TotalRequests;
    } else {
        result.TotalResponseTimeAverage = 0;
    }

    return result;
}

bool TPessimizationTracker::IsOverPessimized() const noexcept {
    return OverPessimized_;
}

void TPessimizationTracker::Dump(NJson::TJsonWriter& out) const noexcept {
    out.OpenArray("backends");
    for (const auto& entry : Backends_) {
        auto& backend = *entry.second.Get();
        out.OpenMap();
        backend->PrintProxyInfo(out);
        backend->Dump(out, true);
        out.CloseMap();
    }
    out.CloseArray();
}
