#include "cpu_usage.h"

#include <balancer/kernel/helpers/syscalls.h>
#include <library/cpp/coroutine/engine/impl.h>

#ifdef _unix_
#   include <sys/resource.h>
#   include <sys/time.h>
#   ifdef _darwin_
#       include <mach/mach.h>
#   endif
#endif

namespace NSrvKernel {
    namespace {
#if defined(_linux_)
        TErrorOr<TDuration> GetCpuUsage() {
            struct rusage cpu = {};
            Y_PROPAGATE_ERROR(Y_SYSCALL(getrusage(RUSAGE_THREAD, &cpu)));
            return TDuration(cpu.ru_utime) + TDuration(cpu.ru_stime);
        }
#elif defined(_darwin_)
        // https://stackoverflow.com/a/55505972/937357
        // No guarantees it even works
        int getrusage_thread(struct rusage *rusage) {
            int ret = -1;
            thread_basic_info_data_t info = {};
            mach_msg_type_number_t info_count = THREAD_BASIC_INFO_COUNT;
            kern_return_t kern_err = thread_info(
                mach_thread_self(),
                THREAD_BASIC_INFO,
                (thread_info_t)&info,
                &info_count
            );
            if (kern_err == KERN_SUCCESS) {
                rusage->ru_utime.tv_sec = info.user_time.seconds;
                rusage->ru_utime.tv_usec = info.user_time.microseconds;
                rusage->ru_stime.tv_sec = info.system_time.seconds;
                rusage->ru_stime.tv_usec = info.system_time.microseconds;
                ret = 0;
            } else {
                errno = EINVAL;
            }
            return ret;
        }

        TErrorOr<TDuration> GetCpuUsage() {
            struct rusage cpu = {};
            Y_PROPAGATE_ERROR(Y_SYSCALL(getrusage_thread(&cpu)));
            return TDuration(cpu.ru_utime) + TDuration(cpu.ru_stime);
        }
#else
        TErrorOr<TDuration> GetCpuUsage() {
            return TDuration();
        }
#endif
     }

    TCoroutine StartCpuAndTimeMeasurer(
        TContExecutor& exec,
        ICpuMeasureCallback* measureCallback1,
        ICpuMeasureCallback* measureCallback2
    ) noexcept {
        static const TDuration period = TDuration::Seconds(1);

        return TCoroutine(ECoroType::Service, "cpu_usage_measure_cont", &exec,
            [&exec, measureCallback1, measureCallback2] {
            Y_TRY(TError, err) {
                TInstant time0 = TInstant::Now();
                TDuration cpu0;
                Y_PROPAGATE_ERROR(GetCpuUsage().AssignTo(cpu0));

                for (;;) {
                    if (auto err = exec.Running()->SleepT(period); err != ETIMEDOUT) {
                        return Y_MAKE_ERROR(TSystemError(err));
                    }

                    TInstant time1 = TInstant::Now();
                    TDuration cpu1;
                    Y_PROPAGATE_ERROR(GetCpuUsage().AssignTo(cpu1));

                    TCpuAndTimeDelta delta{
                        .TotalUsage = cpu1 - cpu0,
                        .MeasureTime = time1 - time0
                    };

                    if (measureCallback1) {
                        measureCallback1->OnMeasure(delta);
                    }
                    if (measureCallback2) {
                        measureCallback2->OnMeasure(delta);
                    }

                    time0 = time1;
                    cpu0 = cpu1;
                }

                return {};
            } Y_CATCH {
                if (!ErrorHasErrno(err, {ECANCELED})) {
                    err.Throw(); // let it die
                }
            }
        });
    }
 }
