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

    File: FutexBarrier.hpp
    Date: 2023-09-20
    Author: Reece
***/
#pragma once

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

namespace Aurora::Threading::Waitables
{
    struct FutexBarrier final
    {
        inline constexpr FutexBarrier(AuUInt32 uExpectThreads) :
            uAtomicSleeping(uExpectThreads),
            uAtomicState(0)
        {}

        AU_NO_COPY_NO_MOVE(FutexBarrier);

        inline bool Enter()
        {
            bool bTimeout {};
            return this->Park(bTimeout, true, 0);
        }
        
        inline AuPair<bool /* bSuccess */, bool /* bTimeOut */>
            EnterTimed(AuUInt64 uTimeoutAbsNS)
        {
            bool bTimeout {};
            bool bSuccess = this->Park(bTimeout, true, uTimeoutAbsNS);
            return AuMakePair(bSuccess, bTimeout);
        }

        inline AuPair<bool /* bSuccess */, bool /* bTimeOut */>
            EnterTimedEx(AuUInt64 uTimeoutAbsNS,
                         AuVoidFunc callOnce)
        {
            bool bTimeout {}, bSuccess {};
            bSuccess = this->Park(bTimeout, 
                                  true,
                                  uTimeoutAbsNS,
                                  callOnce);
            return AuMakePair(bSuccess, bTimeout);
        }

        inline bool TryEnter()
        {
            return this->TryEnterNoSpin();
        }

        inline bool TryWait()
        {
            return this->TryChkNoSpin();
        }

        inline bool Wait(AuUInt64 uTimeoutAbsNS = 0)
        {
            bool bTimeout {};
            (void)this->Park(bTimeout, false, uTimeoutAbsNS);
            return !bTimeout;
        }

    private:

        auline bool Park(bool &bTimedOut,
                         bool bTakeOwnership,
                         AuUInt64 uTimeoutAbsNS,
                         AuVoidFunc callOnce = {})
        {
            static const AuUInt32 kRef { 0 };
            
            bTimedOut = false;
            
            if (bTakeOwnership)
            {
                if (!this->TryEnterNoSpin(callOnce))
                {
                    return false;
                }
            }
            
            if (this->TryChkNoSpin())
            {
                return true;
            }

            (void)TryWaitOnAddress((const void *)&this->uAtomicState, &kRef, sizeof(kRef));
            
            while (!this->TryChkNoSpin())
            {
                if (!WaitOnAddressSteady((const void *)&this->uAtomicState, &kRef, sizeof(kRef), uTimeoutAbsNS))
                {
                    if (this->TryChkNoSpin())
                    {
                        return true;
                    }
                    else
                    {
                        bTimedOut = true;
                        return false;
                    }
                }
            }
            
            return true;
        }
        
        auline bool TryEnterNoSpin(AuVoidFunc callOnce = {})
        {
            AuUInt32 uState;
                
            while ((uState = AuAtomicLoad(&this->uAtomicSleeping)) != 0)
            {
                if (AuAtomicCompareExchange(&this->uAtomicSleeping, uState - 1, uState) == uState)
                {
                    if (uState == 1)
                    {
                        if (callOnce)
                        {
                            callOnce();
                        }

                        AuAtomicStore(&this->uAtomicState, 1u);
                        WakeAllOnAddress((const void *)&this->uAtomicState);
                    }
                    
                    return true;
                }
            }
            
            return false;
        }

        auline bool TryChkNoSpin()
        {
            return AuAtomicLoad(&this->uAtomicState) == 1;
        }

    public:
        AuAUInt32 uAtomicState {};
        AuAUInt32 uAtomicSleeping {};
    };
}