#pragma once

#include <balancer/kernel/helpers/cast.h>
#include <balancer/kernel/helpers/helper_macro.h>

#include <util/datetime/base.h>
#include <util/generic/maybe.h>
#include <util/generic/ymath.h>

#include <algorithm>
#include <array>
#include <limits>
#include <cmath>

namespace NSrvKernel::NRpsLimiter {

    using TUi128 = unsigned __int128;

    // See https://blog.cloudflare.com/counting-things-a-lot-of-different-things/

    class [[nodiscard]] TLocalRate {
    public:
        static TMaybe<TLocalRate> Create(TDuration window, TInstant start={}) noexcept {
            if (!window) {
                return Nothing();
            }
            TLocalRate res;
            res.Window_ = window;
            res.Start_ = start;
            return res;
        }

        bool IncCnt(float cnt) noexcept {
            if (!IsFinite(cnt)) {
                return false;
            }
            Curr_ = std::clamp(Curr_ + cnt, -Max<float>(), Max<float>());
            return true;
        }

        void IncTime(TDuration dt) noexcept {
            if (const auto delta = Epoch(DT_ + dt) - Epoch(DT_)) {
                Prev_ = (delta == 1) ? Curr_ : 0;
                Curr_ = 0;
            }
            DT_ += dt;
        }

        [[nodiscard]]
        double Rate() const noexcept {
            return Curr_ + Prev_ * PrevW();
        }

        [[nodiscard]]
        TUi128 Pack() const noexcept {
            TUi128 s = 0;
            s |= DT_.MicroSeconds();
            s <<= 32;
            s |= CopyCast<ui32>(Curr_);
            s <<= 32;
            s |= CopyCast<ui32>(Prev_);
            return s;
        }

        void Unpack(TUi128 s) noexcept {
            Prev_ = CopyCast<float>(ui32(s));
            s >>= 32;
            Curr_ = CopyCast<float>(ui32(s));
            s >>= 32;
            DT_ = TDuration::MicroSeconds(ui64(s));
        }

        TInstant Time() const noexcept {
            return Start_ + DT_;
        }

        Y_BALANCER_TUPLE_EQ(TLocalRate, Window_, Start_, DT_, Prev_, Curr_)
        Y_BALANCER_GETTER(Window)
        Y_BALANCER_GETTER(Start)
        Y_BALANCER_GETTER(DT)
        Y_BALANCER_GETTER(Prev)
        Y_BALANCER_GETTER(Curr)

    private:
        TLocalRate() = default;

        [[nodiscard]]
        ui64 Epoch(TDuration dt) const noexcept {
            return (dt.MicroSeconds() + Window_.MicroSeconds() - 1) / Window_.MicroSeconds();
        }
        [[nodiscard]]
        double PrevW() const noexcept {
            auto w = double(DT_.MicroSeconds()) / Window_.MicroSeconds();
            return std::ceil(w) - w;
        }

    private:
        TDuration Window_;
        TInstant Start_;
        TDuration DT_;
        float Prev_ = 0;
        float Curr_ = 0;
    };


    class [[nodiscard]] TPeerRate {
    public:
        static TMaybe<TPeerRate> Create(TDuration window, double rate=0) noexcept {
            if (!window || !IsFinite(rate)) {
                return Nothing();
            }
            TPeerRate res;
            res.Window_ = window;
            res.Rate_ = rate;
            return res;
        }

        void Update(TPeerRate r) noexcept {
            if (Window_ != r.Window_) {
                r.Rate_ = r.Rate_ * (Window_ / r.Window_);
            }
            Rate_ = r.Rate_;
        }

        [[nodiscard]]
        ui64 Pack() const noexcept {
            return CopyCast<ui64>(Rate_);
        }

        void Unpack(ui64 s) noexcept {
            Rate_ = CopyCast<double>(s);
        }

        Y_BALANCER_TUPLE_EQ(TPeerRate, Window_, Rate_);
        Y_BALANCER_GETTER(Window)
        Y_BALANCER_GETTER(Rate)

    private:
        TPeerRate() = default;

    private:
        TDuration Window_;
        double Rate_ = 0;
    };


    struct [[nodiscard]] TPeerTime {
        TInstant RemoteTime;
        TInstant LocalTime;

        Y_BALANCER_TUPLE_EQ(TPeerTime, RemoteTime, LocalTime)

    public:
        TUi128 Pack() const noexcept {
            TUi128 s = 0;
            s |= LocalTime.MicroSeconds();
            s <<= 64;
            s |= RemoteTime.MicroSeconds();
            return s;
        }

        void Unpack(TUi128 s) noexcept {
            RemoteTime = TInstant::MicroSeconds(ui64(s));
            s >>= 64;
            LocalTime = TInstant::MicroSeconds(ui64(s));
        }
    };


    template <class TUi>
    class TAtomicStorage : TNonCopyable {
        static_assert(std::is_integral_v<TUi> && std::is_unsigned_v<TUi>);
    public:
        explicit TAtomicStorage(TUi s = 0) noexcept
            : Storage_(s)
        {}

        template <class TUpdater>
        void Update(TUpdater&& updater) noexcept {
            static_assert(std::is_same_v<decltype(updater(TUi())), TUi>);
            TUi prev = Storage_.load();
            while (true) {
                const auto next = updater(prev);
                if (next == prev) {
                    break;
                }
                if (Storage_.compare_exchange_weak(prev, next)) {
                    break;
                }
            }
        }

        TUi Get() const noexcept {
            return Storage_.load();
        }

    private:
        std::atomic<TUi> Storage_ = {};
    };


    template <class TUi>
    class alignas(64) TAlignedAtomicStorage : public TAtomicStorage<TUi> {
    public:
        using TAtomicStorage<TUi>::TAtomicStorage;
    };
}
