#pragma once

#include <balancer/kernel/ctl/children_process_common.h>
#include <balancer/kernel/ctl/children_process_common.h_serialized.h>

#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/cpu/cpu_usage.h>
#include <balancer/kernel/stats/manager.h>

#include <util/generic/serialized_enum.h>

namespace NSrvKernel {

    namespace NDns {
        struct TStatsCounters;
    }

    struct alignas(64) TThreadStatus : TNonCopyable {
    public:
        void UpdateStatus() noexcept {
            AtomicSet(Status_, 1);
        }

        [[nodiscard]] bool GetStatusAndReset() noexcept {
            return AtomicSwap(&Status_, 0);
        }

    private:
        TAtomic Status_ = 0;
    };

    class TWorkerCpuStat final : public ICpuMeasureCallback, public NCoro::IEnterPollerCallback {
    public:
        struct TStats {
            TSharedCounter ContsReady;
            TSharedCounter ContsScheduled;
            TSharedCounter ContsWaiting;
            TSharedCounter ContsAllocReleasedSize;
            TSharedCounter ContsAllocNotReleasedSize;
            TSharedCounter ContsAllocated;
            TSharedCounter ContsCycleCounter;
            TSharedCounter DesyncCounter;
            TSharedHistogram CpuUsageHist;
            TSharedHistogram ContsCycleHist;
        };

        TWorkerCpuStat(TStats stats, TThreadStatus* status)
            : Stats_(std::move(stats))
            , Status_(status)
        {}

    private:
        void OnMeasure(TCpuAndTimeDelta delta) noexcept {
            if (delta.MeasureTime.MilliSeconds() > 0) {
                ui8 load = 100.0 * delta.TotalUsage / delta.MeasureTime;
                Stats_.CpuUsageHist.AddValue(load);
            }

            TContExecutor* executor = RunningCont()->Executor();
            Stats_.ContsReady.Set(executor->TotalReadyConts());
            Stats_.ContsScheduled.Set(executor->TotalScheduledConts());
            Stats_.ContsWaiting.Set(executor->TotalWaitingConts());
            auto allocStats = executor->GetAllocatorStats();
            Stats_.ContsAllocReleasedSize.Set(allocStats.ReleasedSize);
            Stats_.ContsAllocNotReleasedSize.Set(allocStats.NotReleasedSize);
            Stats_.ContsAllocated.Set(allocStats.NumOfAllocated);

            if (Status_) {
                Status_->UpdateStatus();
            }
        }

        void OnEnterPoller() {
            if (ExitPollerTime_) {
                UpdateCycle(Now() - ExitPollerTime_);
                ExitPollerTime_ = TInstant::Zero();
            }
        }

        void OnExitPoller() {
            ExitPollerTime_ = Now();
        }

        void UpdateCycle(TDuration cycleDuration) noexcept {
            Stats_.ContsCycleCounter.Inc();
            Stats_.ContsCycleHist.AddValue(cycleDuration.MicroSeconds());

            if (cycleDuration > TDuration::MilliSeconds(10)) {
                Stats_.DesyncCounter.Inc();
            }
        }

        TStats Stats_;

        TInstant ExitPollerTime_;

        TThreadStatus* Status_ = nullptr;
    };

    class TProcessStat {
    public:
        TProcessStat(TSharedStatsManager& statsManager, TVector<TThreadStatus>& statuses, TString processType);

        THolder<TWorkerCpuStat> GetWorkerCpuStat(size_t workerId);

        THolder<NDns::TStatsCounters> GetDnsCounters(size_t workerId);

    private:
        TString ProcessName_;
        TWorkerCpuStat::TStats Stats_;

        THolder<NDns::TStatsCounters> DnsCounters_;
        TVector<TThreadStatus>& ThreadsStatuses_;
    };


    class TProcessStatOwner {
    public:
        TProcessStatOwner(TSharedStatsManager& statsManager, size_t countOfChildren)
            : StatsManager_(&statsManager)
            , MasterStats_(MakeHolder<TProcessStat>(statsManager, ThreadsStatuses_, "master"))
            , ThreadsStatuses_(countOfChildren + 1)
        {}

        void RegisterStatForProcess(NProcessCore::TChildProcessMask mask) {
            for (const auto& type : GetEnumAllValues<NProcessCore::TChildProcessType>()) {
                if (mask.HasFlags(type)) {
                    const size_t index = ChildProcessTypeToIndex(type);
                    Y_VERIFY(!StatByProcessType_[index], "process type has already been registered");
                    StatByProcessType_[index] = MakeHolder<TProcessStat>(*StatsManager_, ThreadsStatuses_, ::ToString(type));
                }
            }
        }

        TProcessStat* GetMasterStat() {
            return MasterStats_.Get();
        }

        TProcessStat* GetProcessStat(NProcessCore::TChildProcessType type) {
            auto* res = StatByProcessType_[ChildProcessTypeToIndex(type)].Get();
            Y_VERIFY(res, "GetProcessStat for non-existent type");
            return res;
        }

        void GetThreadsStatuses(TVector<bool>& statuses) noexcept {
            statuses.resize(ThreadsStatuses_.size());
            for (size_t i = 0ul; i < ThreadsStatuses_.size(); ++i) {
                statuses[i] = ThreadsStatuses_[i].GetStatusAndReset();
            }
        }

    private:
        TSharedStatsManager* const StatsManager_;
        THolder<TProcessStat> MasterStats_;
        std::array<THolder<TProcessStat>, GetEnumItemsCount<NProcessCore::TChildProcessType>()> StatByProcessType_;
        TVector<TThreadStatus> ThreadsStatuses_;
    };
}
