#include "quota.h"

#include <balancer/kernel/rpslimiter/proto/state.pb.h>
#include <balancer/kernel/tvm/tvm.h>

#include <util/generic/algorithm.h>
#include <util/stream/mem.h>
#include <util/string/join.h>

namespace NSrvKernel::NRpsLimiter {


    TErrorOr<TPeerQuotas> ParsePeerQuotas(TStringBuf s) noexcept {
        TQuotaStorageState st;
        try {
            TMemoryInput in(s.data(), s.size());
            st.Load(&in);
        } Y_TRY_STORE(yexception);
        TPeerQuotas hq;
        hq.Name = st.GetHostId();
        hq.Time = TInstant::MicroSeconds(st.GetTstampMicrosec());
        hq.Quotas.reserve(st.QuotaStatesSize());
        for (auto&& q : st.GetQuotaStates()) {
            hq.Quotas.emplace_back(TPeerQuota{
                .Name=q.GetName(),
                .Window=TDuration::MicroSeconds(q.GetIntervalMicrosec()),
                .Rate=q.GetState().GetRate(),
                .IsByQuota=q.GetIsByQuota()
            });
        }
        return hq;
    }

    TString RenderPeerQuotas(const TPeerQuotas& hq) noexcept {
        TQuotaStorageState st;
        st.SetHostId(hq.Name);
        st.SetTstampMicrosec(hq.Time.MicroSeconds());
        st.MutableQuotaStates()->Reserve(hq.Quotas.size());
        for (auto&& q : hq.Quotas) {
            auto* qSt = st.AddQuotaStates();
            qSt->SetName(q.Name);
            qSt->SetIntervalMicrosec(q.Window.MicroSeconds());
            qSt->MutableState()->SetRate(q.Rate);
            qSt->SetIsByQuota(q.IsByQuota);
        }

        TStringBuilder sb;
        st.Save(&sb.Out);
        return sb;
    }


    TQuotaState::TQuotaState(TString selfName, TInstant now, TVector<TString> peerNames, TVector<TQuota> quotas)
        : SelfName_(selfName)
        , Now_(now)
        , ByQuotasLock_(std::make_unique<std::mutex>())
    {
        Sort(peerNames);
        Sort(quotas);
        Erase(peerNames, selfName);
        LocalQuotas_.resize(quotas.size());

        PeerQuotas_.resize(peerNames.size());
        PeerTime_.resize(peerNames.size());
        for (auto&& h : PeerQuotas_) {
            h.resize(quotas.size());
        }
        PeerByQuotas_.resize(peerNames.size());
        PeerByQuotasLock_.resize(peerNames.size());
        for (size_t idx : xrange(peerNames.size())) {
            PeerByQuotasLock_[idx] = std::make_unique<std::mutex>();
        }

        Quotas_.insert(Quotas_.end(), quotas.begin(), quotas.end());
        for (const auto [idx, h] : Enumerate(peerNames)) {
            Y_ENSURE_EX(PeerIdxByName_.emplace(h, idx).second,
                NConfig::TConfigParseError() << " host name is not unique " << h.Quote());
        }
        for (const auto [idx, q] : Enumerate(quotas)) {
            Y_ENSURE_EX(q.Window >= TDuration::MilliSeconds(100),
                NConfig::TConfigParseError() << " too small window for quota " << TString(q.Name).Quote());
            Y_ENSURE_EX(QuotaIdxByName_.emplace(q.Name, idx).second,
                NConfig::TConfigParseError() << " quota name is not unique " << TString(q.Name).Quote());
            for (const auto& byQuotaCfg : q.By) {
                if (byQuotaCfg.data_type() == "tvm-service") {
                    // Just initialization of client object
                    TClientMap::Instance().GetClient(IntFromString<NTvmAuth::TTvmId, 10>(byQuotaCfg.value()));
                }
            }
        }
    }

    TPeerQuotas TQuotaState::GenPeerQuotasFromLocal(TInstant now) const noexcept {
        TPeerQuotas res;
        res.Name = SelfName_;
        res.Time = now;
        res.Quotas.reserve(Quotas_.size());
        for (size_t idx : xrange(Quotas_.size())) {
            auto rate = LocalQuotaRate(idx, now);
            if (!rate) {
                continue;
            }
            const auto& quota = Quotas_[idx];
            res.Quotas.emplace_back(TPeerQuota{
                .Name=quota.Name,
                .Window=quota.Window,
                .Rate=rate
            });
        }

        TVector<TByQuota> byQuotas;
        {
            std::scoped_lock guard{*ByQuotasLock_};
            for (const auto& it : ByQuotas_) {
                std::scoped_lock quotaGuard{*it.second->Lock};
                byQuotas.push_back(*it.second);
            }
        }
        for (const auto& byQuota : byQuotas) {
            const auto& quota = byQuota.Quota;
            double rate = QuotaRate(quota, byQuota.Rate, now);
            if (!rate) {
                std::scoped_lock guard{*ByQuotasLock_};
                const auto it = ByQuotas_.find(quota.Name);
                if (it != ByQuotas_.end()) {
                    ByQuotas_.erase(it);
                }
                continue;
            }
            res.Quotas.emplace_back(TPeerQuota{
                .Name=quota.Name,
                .Window=quota.Window,
                .Rate=rate,
                .IsByQuota=true
            });
        }

        return res;
    }

    bool TQuotaState::UpdatePeerQuotas(const TPeerQuotas& quotas, TInstant now) noexcept {
        if (quotas.Name.empty() || SelfName_ == quotas.Name) {
            return false;
        }
        const auto peerIdx = PeerIdx(quotas.Name);
        if (!peerIdx) {
            return false;
        }
        bool old = false;
        PeerTime_[*peerIdx].Update([&](auto s) {
            TPeerTime pt;
            pt.Unpack(s);
            old = (pt.RemoteTime >= quotas.Time);
            if (old) {
                return s;
            }
            pt.RemoteTime = quotas.Time;
            pt.LocalTime = now;
            return pt.Pack();
        });
        if (old) {
            return false;
        }

        auto& hq = PeerQuotas_[*peerIdx];
        auto& peerByQuotas = PeerByQuotas_[*peerIdx];
        auto& peerByQuotasLock = PeerByQuotasLock_[*peerIdx];

        static TInstant cleanupTime = TInstant::Zero();
        const TDuration cleanupInterval = TDuration::Seconds(1);
        if ((now - cleanupTime) >= cleanupInterval) {
            cleanupTime = now;
            std::scoped_lock guard(*peerByQuotasLock);
            auto it = peerByQuotas.begin();
            while (it != peerByQuotas.end()) {
                const auto& byQuota = it->second;
                if (!byQuota.Rate) {
                    it = peerByQuotas.erase(it);
                    continue;
                }
                double rate = QuotaRate(byQuota.Quota, byQuota.Rate, now);
                if (!rate) {
                    it = peerByQuotas.erase(it);
                } else {
                    ++it;
                }
            }
        }

        bool updated = false;
        for (auto&& q : quotas.Quotas) {
            TString name = q.Name;
            if (q.IsByQuota) {
                name = ParseByQuotaName(q.Name)[0];
            }
            auto idx = QuotaIdx(name);
            if (!idx) {
                continue;
            }
            auto upd = TPeerRate::Create(q.Window, q.Rate);
            if (!upd) {
                continue;
            }
            const auto& quota = Quotas_[*idx];
            auto rate = *TPeerRate::Create(quota.Window);
            if (!q.IsByQuota) {
                hq[*idx].Update([&](auto s) {
                    rate.Unpack(s);
                    rate.Update(*upd);
                    return rate.Pack();
                });
            } else {
                std::scoped_lock guard(*peerByQuotasLock);
                TByQuota byQuota{quota, q.Name};
                rate.Unpack(byQuota.Rate);
                rate.Update(*upd);
                byQuota.Rate = rate.Pack();
                peerByQuotas[q.Name] = byQuota;
            }
            updated = true;
        }
        return updated;
    }
}

using namespace NSrvKernel;

template <>
void Out<NRpsLimiter::TQuota>(IOutputStream& out, const NRpsLimiter::TQuota& q) {
    out << "{.Name=" << q.Name.Quote()
        << ",.Start=TInstant::MicroSeconds(" << q.Start.MicroSeconds() << ")"
        << ",.Window=TDuration::MicroSeconds(" << q.Window.MicroSeconds() << ")"
        << ",.Limit=" << q.Limit << "}";
}

template <>
void Out<NRpsLimiter::TPeerQuota>(IOutputStream& out, const NRpsLimiter::TPeerQuota& q) {
    out << "{.Name=" << q.Name.Quote()
        << ",.Window=TDuration::MicroSeconds(" << q.Window.MicroSeconds() << ")"
        << ",.Rate=" << q.Rate << "}";
}

template <>
void Out<NRpsLimiter::TPeerQuotas>(IOutputStream& out, const NRpsLimiter::TPeerQuotas& q) {
    out << "{.Name=" << q.Name.Quote()
        << ",.Time=TInstant::MicroSeconds(" << q.Time.MicroSeconds() << ")"
        << ",.Quotas={" << JoinSeq(",", q.Quotas) << "}}";
}
