#include "threadpool.h"
#include <library/cpp/testing/unittest/registar.h>
#include <util/system/event.h>

#if defined(_linux_)
#include <pthread.h>
#include <sched.h>
#include <errno.h>
#include <sys/resource.h>
#endif

Y_UNIT_TEST_SUITE(TSaasPrioThreadTest) {
    static bool SupportedPlatform() {
#if defined(_linux_)
        return true;
#else
        return false;
#endif
    }

    static bool CheckSchedParam(bool& isCpuIdle, int& prio) {
#if defined(_linux_)
        int rc;
        struct sched_param param;
        int cpuScheduler;
        pthread_t tid;
        prio = getpriority(PRIO_PROCESS, 0);
        if (errno)
            return false;

        tid = pthread_self();
        rc = pthread_getschedparam(tid, &cpuScheduler, &param);
        if (rc)
            return false;

        isCpuIdle = 0 != (cpuScheduler & SCHED_BATCH);
        return true;
#else
        Y_UNUSED(isCpuIdle)
        Y_UNUSED(prio);
        return false;
#endif
    }

    void DoTest(TScheduling opts, int cpuNice) {
        auto tp = MakeHolder<TPriorityThreadPool>(opts, cpuNice);
        auto queue = MakeHolder<TThreadPool>(tp.Get());
        queue->Start(1, 1);
        bool result = false;
        TAutoEvent completed;

        int expectedPrio = getpriority(PRIO_PROCESS, 0) + cpuNice;
        bool added = queue->AddFunc([&] {
            bool isCpuIdle;
            int prio;
            result = CheckSchedParam(isCpuIdle, prio) && isCpuIdle == !!(opts & TScheduling::CpuBatch) && prio >= expectedPrio;
            completed.Signal();
        });
        UNIT_ASSERT(added);
        UNIT_ASSERT(completed.WaitT(YaIsDebuggerPresent() ? TDuration::Days(1) : TDuration::Seconds(1)));
        queue->Stop();
        UNIT_ASSERT(result);
    }

    Y_UNIT_TEST(THasPrio) {
        if (!SupportedPlatform())
            return;

        DoTest(TScheduling(), 1);
    }

    Y_UNIT_TEST(THasScheduling) {
        if (!SupportedPlatform())
            return;

        DoTest(TScheduling::CpuBatch, 0);
    }
}
