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

    File: InitOnce.hpp
    Date: 2023-09-17
    Author: Reece
***/
#pragma once

namespace Aurora::Threading
{
    struct InitOnceLocker;
}

namespace __audetail
{
    struct InitOnceABI
    {
        inline bool IsUninitialized();

        inline bool IsReady();

        inline bool TrySet();

        inline void Wait();

    protected:

        inline bool _TryAcquire();

        inline bool _TryAcquireFailable();

        inline void _Finish();

        inline void _Wakeup();

        inline void _FinishFailed();

        inline bool _LockAbsNS(AuUInt64 qwAbsTimeoutInNs /* = 0, infinity*/);

    protected:

        AuAUInt32 uToken_ {};
        AuAUInt32 uSleepers_ {};

        friend struct Aurora::Threading::InitOnceLocker;
    };
}

namespace Aurora::Threading
{
    struct InitOnce final :
        private IWaitable,
        private __audetail::InitOnceABI
    {
        inline bool IsUninitialized()
        {
            return InitOnceABI::IsUninitialized();
        }

        inline bool IsReady()
        {
            return InitOnceABI::IsReady();
        }

        inline bool TrySet()
        {
            return InitOnceABI::TrySet();
        }

        template <typename Callable>
        bool TryInit(const Callable &callback)
        {
            if (InitOnceABI::_TryAcquire())
            {
                callback();
                InitOnceABI::_Finish();
                return true;
            }
            else
            {
                return false;
            }
        }

        template <typename Callable>
        void Call(const Callable &callback)
        {
            if (this->TryInit(callback))
            {
                return;
            }

            if (this->IsReady())
            {
                return;
            }

            this->Wait();
        }

        template <typename Callable>
        bool TryCall(const Callable &callback, bool bWait = true)
        {
            bool bRet {};

            if (InitOnceABI::_TryAcquireFailable())
            {
                if (callback())
                {
                    InitOnceABI::_Finish();
                    return true;
                }
                else
                {
                    InitOnceABI::_FinishFailed();
                    return false;
                }
            }

            if ((!(bRet = this->IsReady())) &&
                (bWait))
            {
                this->Wait();
                bRet = this->IsReady();
            }

            return bRet;
        }

        inline void Wait()
        {
            InitOnceABI::Wait();
        }

        inline IWaitable *ToBarrier()
        {
            return this;
        }

    protected:

        inline bool HasOSHandle(AuMach &mach) override
        {
            return false;
        }

        inline bool HasLockImplementation() override
        {
            return true;
        }

        inline void Lock() override
        {
            this->LockAbsNS(0);
        }

        inline bool LockNS(AuUInt64 qwTimeoutInNs) override
        {
            if (this->IsReady())
            {
                return true;
            }

            return IWaitable::LockNS(qwTimeoutInNs);
        }

        inline bool LockAbsNS(AuUInt64 qwAbsTimeoutInNs /* = 0, infinity*/) override
        {
            return this->_LockAbsNS(qwAbsTimeoutInNs);
        }

        inline bool TryLock() override
        {
            return this->IsReady();
        }

        inline void Unlock() override
        {

        }

    private:
        friend struct InitOnceLocker;
    };

    struct InitOnceLocker
    {
        cstatic bool TryLock(InitOnce *pInitOnce, bool bFailable = false)
        {
            if (bFailable)
            {
                return pInitOnce->_TryAcquireFailable();
            }
            else
            {
                return pInitOnce->_TryAcquire();
            }
        }

        cstatic void Finish(InitOnce *pInitOnce, bool bFailed = false)
        {
            if (bFailed)
            {
                pInitOnce->_FinishFailed();
            }
            else
            {
                pInitOnce->_Finish();
            }
        }

        cstatic bool TryLock(__audetail::InitOnceABI *pInitOnce, bool bFailable = false)
        {
            if (bFailable)
            {
                return pInitOnce->_TryAcquireFailable();
            }
            else
            {
                return pInitOnce->_TryAcquire();
            }
        }

        cstatic void Finish(__audetail::InitOnceABI *pInitOnce, bool bFailed = false)
        {
            if (bFailed)
            {
                pInitOnce->_FinishFailed();
            }
            else
            {
                pInitOnce->_Finish();
            }
        }

        /**
         * Example usage:
         *
         *  static AuInitOnceSmall gSmall; // or static AuInitOnce gInitOnce;
         *  static_assert(sizeof(gSmall) == 8);
         * 
         *  if (AuThreading::InitOnceLocker::TryLock(&gSmall))
         *  {
         *      // My non-AuInitOnce::TryInit/Call() callback logic here
         *      AuThreading::InitOnceLocker::Finish(&gSmall);
         *  }
         *  else
         *  {
         *      gSmall.Wait();
         *  }
         */
    };
}

namespace __audetail
{
    bool InitOnceABI::IsUninitialized()
    {
        return (AuAtomicLoad(&this->uToken_) & 1) == 0;
    }

    bool InitOnceABI::IsReady()
    {
        return (AuAtomicLoad(&this->uToken_) & 2) != 0;
    }

    bool InitOnceABI::TrySet()
    {
        if (this->_TryAcquire())
        {
            this->_Finish();
            return true;
        }
        else
        {
            return false;
        }
    }

    void InitOnceABI::Wait()
    {
        this->_LockAbsNS(0);
    }

    bool InitOnceABI::_TryAcquire()
    {
        return AuAtomicTestAndSet(&this->uToken_, 0) == 0;
    }

    bool InitOnceABI::_TryAcquireFailable()
    {
        if (AuAtomicTestAndSet(&this->uToken_, 0) == 0)
        {
            AuAtomicUnset(&this->uToken_, 2);
            return true;
        }
        else
        {
            return false;
        }
    }
    
    void InitOnceABI::_Finish()
    {
        AuAtomicSet(&this->uToken_, 1);
        this->_Wakeup();
    }

    void InitOnceABI::_Wakeup()
    {
        if (auto uSleepers = AuAtomicLoad(&this->uSleepers_))
        {
            Aurora::Threading::WakeNOnAddress((const void *)&this->uToken_, uSleepers);
        }
    }

    void InitOnceABI::_FinishFailed()
    {
        AuAtomicUnset(&this->uToken_, 0);
        AuAtomicSet(&this->uToken_, 2);
        InitOnceABI::_Wakeup();
    }

    bool InitOnceABI::_LockAbsNS(AuUInt64 qwAbsTimeoutInNs /* = 0, infinity*/)
    {
        auto uCurrent = AuAtomicLoad(&this->uToken_);

        while ((uCurrent & (2 | 4)) == 0)
        {
            AuAtomicAdd(&this->uSleepers_, 1u);
            bool bRet = Aurora::Threading::WaitOnAddressSteady((const void *)&this->uToken_, &uCurrent, sizeof(uCurrent), qwAbsTimeoutInNs);
            AuAtomicSub(&this->uSleepers_, 1u);

            if (!bRet)
            {
                return this->IsReady();
            }

            uCurrent = AuAtomicLoad(&this->uToken_);
        }

        return true;
    }
}