/***
    Copyright (C) 2023 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: FutexCondWaitable.hpp
    Date: 2023-08-19
    Author: Reece
***/
#pragma once

#include "../WakeOnAddress.hpp"
#include "../SpinTime.hpp"

namespace Aurora::Threading::Waitables
{
    struct FutexCondWaitable final 
    {
        inline constexpr FutexCondWaitable()
        {}

        AU_NO_COPY_NO_MOVE(FutexCondWaitable);

        template <class T>
        inline bool WaitNS(T pWaitable, AuUInt64 uRelativeNanoseconds)
        {
            AuAtomicAdd(&this->uAtomicSleeping, 1u);

            if (pWaitable)
            {
                pWaitable->Unlock();
            }

            auto bSuccess = this->TryLock2() || 
                this->SleepOne(Time::SteadyClockNS() + uRelativeNanoseconds);

            if (!bSuccess)
            {
                auto uWaiters = this->uAtomicSleeping;
                auto uWaitCount = 1;

                while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
                {
                    uWaiters = this->uAtomicSleeping;

                    if (uWaiters == 0)
                    {
                        break;
                    }
                }
            }

            if (pWaitable)
            {
                pWaitable->Lock();
            }

            return bSuccess;
        }

        template <class T>
        inline bool WaitAbsNS(T pWaitable, AuUInt64 uSteadyClockNanoseconds)
        {
            AuAtomicAdd(&this->uAtomicSleeping, 1u);

            if (pWaitable)
            {
                pWaitable->Unlock();
            }

            auto bSuccess = this->TryLock2() || 
                this->SleepOne(uSteadyClockNanoseconds);

            if (!bSuccess)
            {
                auto uWaiters = this->uAtomicSleeping;
                auto uWaitCount = 1;

                while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
                {
                    uWaiters = this->uAtomicSleeping;

                    if (uWaiters == 0)
                    {
                        break;
                    }
                }
            }

            if (pWaitable)
            {
                pWaitable->Lock();
            }

            return bSuccess;
        }

        auline void NotifyOne()
        {
            AuUInt32 uWaitCount {};
            AuUInt32 uWaiters {};

            uWaiters = this->uAtomicSleeping;
            if (uWaiters > 0)
            {
                AuAtomicAdd(&this->uAtomicState, 1u);
                WakeOnAddress((const void *)&this->uAtomicState);
                uWaitCount = 1;
            }

            while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
            {
                uWaiters = this->uAtomicSleeping;

                if (uWaiters == 0)
                {
                    return;
                }
            }
        }

        auline void Broadcast()
        {
            AuUInt32 uWaitCount {};
            AuUInt32 uWaiters {};

            uWaiters = this->uAtomicSleeping;
            if (uWaiters > 0)
            {
                AuAtomicAdd(&this->uAtomicState, uWaiters);
                WakeNOnAddress((const void *)&this->uAtomicState, uWaiters);
                uWaitCount = 1;
            }

            while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
            {
                uWaiters = this->uAtomicSleeping;

                if (uWaiters == 0)
                {
                    return;
                }
            }
        }

        auline void NotifyN(AuUInt32 uThreads)
        {
            AuUInt32 uWaitCount {};
            AuUInt32 uWaiters {};

            uWaiters = this->uAtomicSleeping;
            if (uWaiters > 0)
            {
                AuAtomicAdd(&this->uAtomicState, uThreads);
                WakeNOnAddress((const void *)&this->uAtomicState, uThreads);
                uWaitCount = 1;
            }

            while (AuAtomicCompareExchange(&this->uAtomicSleeping, uWaiters - uWaitCount, uWaiters) != uWaiters)
            {
                uWaiters = this->uAtomicSleeping;

                if (uWaiters == 0)
                {
                    return;
                }
            }
        }

        volatile AuUInt32 uAtomicState {};
        volatile AuUInt32 uAtomicSleeping {};
    private:

        auline bool TryLock2()
        {
            for (AU_ITERATE_N(i, AuUInt(1u) << AuUInt(Threading::GetSpinCountTimeout())))
            {
            #if defined(AURORA_ARCH_X86) || defined(AURORA_ARCH_X64)
                _mm_pause();
            #else
                Threading::ContextYield();
            #endif

                auto old = this->uAtomicState;
                if ((old != 0 && AuAtomicCompareExchange(&this->uAtomicState, old - 1, old) == old))
                {
                    return true;
                }
            }

            return false;
        }
        
        auline bool SleepOne(AuUInt64 qwTimeout)
        {
            static const AuUInt32 kRef { 0 };

            while (!TryLock2())
            {
                bool bStatus {};

                bStatus = WaitOnAddressSteady((void *)&this->uAtomicState, &kRef, sizeof(kRef), qwTimeout);

                if (!bStatus)
                {
                    return TryLock2();
                }
            }

            return true;
        }
    };
}