#pragma once

#include "rate.h"

#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/client_request/port.h>

#include <library/cpp/config/config.h>
#include <library/cpp/iterator/enumerate.h>

#include <util/datetime/base.h>
#include <util/generic/deque.h>
#include <util/generic/hash.h>
#include <util/generic/map.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>
#include <util/generic/xrange.h>
#include <util/string/split.h>

#include <tuple>

#include <balancer/kernel/rpslimiter/rpslimiter_instance.cfgproto.pb.h>

namespace NSrvKernel::NRpsLimiter {

    struct TQuota {
        TString Name;
        TInstant Start;
        TDuration Window;
        double Limit = 0;
        TVector<TByQuotaCfg> By;

        Y_BALANCER_TUPLE_OPS(TQuota, Name, Start, Window, Limit)
    };


    struct TPeerQuota {
        TString Name;
        TDuration Window;
        double Rate = 0;
        bool IsByQuota = false;

        Y_BALANCER_TUPLE_OPS(TPeerQuota, Name, Window, Rate, IsByQuota)
    };


    struct TPeerQuotas {
        TString Name;
        TInstant Time;
        TVector<TPeerQuota> Quotas;

        Y_BALANCER_TUPLE_EQ(TPeerQuotas, Name, Time, Quotas)
    };

    struct TByQuota {
        TQuota Quota;
        TUi128 Rate = 0;
        std::shared_ptr<std::mutex> Lock;

        TByQuota() {}
        TByQuota(const TQuota& quota, const TString& fullName)
            : Quota(quota)
            , Lock(std::make_shared<std::mutex>())
        {
            Quota.Name = fullName;
        }
    };
    using TByQuotaPtr = TAtomicSharedPtr<TByQuota>;

    TErrorOr<TPeerQuotas> ParsePeerQuotas(TStringBuf s) noexcept;

    TString RenderPeerQuotas(const TPeerQuotas& hq) noexcept;


    class TQuotaState : public TMoveOnly {
        using TNameIdx = THashMap<TString, size_t>;

    public:
        TQuotaState(TString selfName, TInstant now, TVector<TString> peerNames, TVector<TQuota> quotas);

        Y_BALANCER_GETTER(SelfName)

        TMaybe<size_t> QuotaIdx(TStringBuf name) const noexcept {
            return Idx(QuotaIdxByName_, name);
        }

        TMaybe<size_t> PeerIdx(TStringBuf name) const noexcept {
            return Idx(PeerIdxByName_, name);
        }

        const TQuota& QuotaInfo(size_t idx) const noexcept {
            return Quotas_[idx];
        }

        double TotalQuotaRate(size_t idx, TInstant now) const noexcept {
            return LocalQuotaRate(idx, now) + PeerQuotaRate(idx, now);
        }

        double PeerQuotaRate(size_t idx, TInstant now) const noexcept {
            const auto& quota = Quotas_[idx];
            double res = 0;
            auto rate = *TPeerRate::Create(quota.Window);
            for (auto peerIdx : xrange(PeerTime_.size())) {
                TPeerTime pt;
                pt.Unpack(PeerTime_[peerIdx].Get());
                if (pt.LocalTime + quota.Window > now) {
                    auto rt = rate;
                    rt.Unpack(PeerQuotas_[peerIdx][idx].Get());
                    res += rt.Rate();
                }
            }
            return res;
        }

        double LocalQuotaRate(size_t idx, TInstant now) const noexcept {
            const auto& quota = Quotas_[idx];
            auto current = LocalQuotas_[idx].Get();
            return QuotaRate(quota, current, now);
        }

        void IncLocalQuota(size_t idx, TInstant now, float cnt) noexcept {
            const auto& quota = Quotas_[idx];
            auto rate = *TLocalRate::Create(quota.Window, quota.Start);
            LocalQuotas_[idx].Update([&](auto s) {
                rate.Unpack(s);
                rate.IncTime(now - rate.Time());
                rate.IncCnt(cnt);
                return rate.Pack();
            });
        }

        TQuota ByQuotaInfo(const TQuota& quota, TStringBuf byValue) noexcept {
            const TString fullName = ByQuotaFullName(quota.Name, byValue);
            return FindOrCreateByQuotaCopy(quota, fullName)->Quota;
        }

        double TotalByQuotaRate(const TQuota& quota, TInstant now) const noexcept {
            return  LocalByQuotaRate(quota, now) + PeerByQuotaRate(quota, now);;
        }

        double LocalByQuotaRate(const TQuota& quota, TInstant now) const noexcept {
            TMaybe<TByQuotaPtr> found;
            {
                std::scoped_lock guard{*ByQuotasLock_};
                found = FindExistingByQuotaUnsafe(quota);
            }
            if (!found) {
                return 0;
            }
            const auto& byQuota = **found;
            std::scoped_lock quotaGuard{*byQuota.Lock};
            return QuotaRate(byQuota.Quota, byQuota.Rate, now);
        }

        double PeerByQuotaRate(const TQuota& quota, TInstant now) const noexcept {
            double res = 0;
            auto rate = *TPeerRate::Create(quota.Window);
            for (auto peerIdx : xrange(PeerTime_.size())) {
                TPeerTime pt;
                pt.Unpack(PeerTime_[peerIdx].Get());
                if (pt.LocalTime + quota.Window > now) {
                    auto rt = rate;
                    {
                        std::scoped_lock guard(*PeerByQuotasLock_[peerIdx]);
                        rt.Unpack(PeerByQuotas_[peerIdx][quota.Name].Rate);
                    }
                    res += rt.Rate();
                }
            }
            return res;
        }

        void IncLocalByQuota(const TQuota& quota, TInstant now, float cnt) noexcept {
            TByQuotaPtr byQuota;
            {
                std::scoped_lock guard{*ByQuotasLock_};
                byQuota = FindOrCreateByQuotaUnsafe(quota, quota.Name);
            }
            std::scoped_lock quotaGuard{*byQuota->Lock};
            auto rate = *TLocalRate::Create(byQuota->Quota.Window, byQuota->Quota.Start);
            rate.Unpack(byQuota->Rate);
            rate.IncTime(now - rate.Time());
            rate.IncCnt(cnt);
            byQuota->Rate = rate.Pack();
        }

        [[nodiscard]]
        bool UpdatePeerQuotas(const TPeerQuotas& quotas, TInstant now) noexcept;

        TPeerQuotas GenPeerQuotasFromLocal(TInstant now) const noexcept;

    private:
        double QuotaRate(const TQuota& quota, TUi128 current, TInstant now) const noexcept {
            auto rate = *TLocalRate::Create(quota.Window, quota.Start);
            rate.Unpack(current);
            rate.IncTime(now - rate.Time());
            return rate.Rate();
        }

        static TMaybe<size_t> Idx(const TNameIdx& byName, TStringBuf name) noexcept {
            if (auto it = byName.find(name); it != byName.end()) {
                return it->second;
            }
            return Nothing();
        }

        TString ByQuotaFullName(const TString& name, TStringBuf byValue) const noexcept {
            return name + ":" + byValue;
        }

        TVector<TString> ParseByQuotaName(const TString& name) const noexcept {
            return StringSplitter(name).Split(':');
        }

        TMaybe<TByQuotaPtr> FindExistingByQuotaUnsafe(const TQuota& quota) const noexcept {
            auto it = ByQuotas_.find(quota.Name);
            if (it != ByQuotas_.end()) {
                return it->second;
            }
            return Nothing();
        }

        TByQuotaPtr FindOrCreateByQuotaUnsafe(const TQuota& quota, const TString& fullName) noexcept {
            auto it = ByQuotas_.find(fullName);
            if (it != ByQuotas_.end()) {
                return it->second;
            }
            const auto insert = ByQuotas_.emplace(
                make_pair(fullName, MakeAtomicShared<TByQuota>(quota, fullName))
            );
            return insert.first->second;
        }

        TByQuotaPtr FindOrCreateByQuotaCopy(const TQuota& quota, const TString& fullName) noexcept {
            std::scoped_lock guard{*ByQuotasLock_};
            return FindOrCreateByQuotaUnsafe(quota, fullName);
        }

    private:
        TString SelfName_;
        TInstant Now_;
        TVector<TQuota> Quotas_;    // Quotas_[quotaIdx]
        TNameIdx QuotaIdxByName_;
        TNameIdx PeerIdxByName_;
        TDeque<TAlignedAtomicStorage<TUi128>> LocalQuotas_;     // LocalQuotas_[quotaIdx]
        TDeque<TAlignedAtomicStorage<TUi128>> PeerTime_;        // PeerTime_[peerIdx]
        TVector<TDeque<TAtomicStorage<ui64>>> PeerQuotas_;      // PeerQuotas_[peerIdx][quotaIdx]

        std::unique_ptr<std::mutex> ByQuotasLock_;
        mutable TMap<TString, TByQuotaPtr> ByQuotas_;

        TVector<std::unique_ptr<std::mutex>> PeerByQuotasLock_;
        mutable TVector<TMap<TString, TByQuota>> PeerByQuotas_;
    };
}
