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

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

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

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

        AU_NO_COPY_NO_MOVE(FutexWaitable);

        inline bool TryLockNoSpin()
        {
            return AuAtomicTestAndSet(&this->uAtomicState, 0u) == 0;
        }

        inline bool TryLock() override
        {
            if (TryLockNoSpin())
            {
                return true;
            }

        #if defined(AURORA_ARCH_X86) || defined(AURORA_ARCH_X64)
            for (AU_ITERATE_N(i, AuUInt(GetTotalSpinCountTimeout())))
            {
                _mm_pause();

                if (TryLockNoSpin())
                {
                    return true;
                }
            }
        #else
            static const AuUInt32 kRef { 1 };
            if (TryWaitOnAddress((void *)&this->uAtomicState, &kRef, sizeof(kRef)))
            {
                if (TryLockNoSpin())
                {
                    return true;
                }
            }
        #endif

            return false;
        }

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

        inline bool HasLockImplementation() override
        {
            return true;
        }

        inline void Unlock() override
        {
        #if defined(AURORA_COMPILER_MSVC)
            this->uAtomicState = 0;
        #else
            __sync_lock_release(&this->uAtomicState);
        #endif

            if (auto uSleeping = this->uAtomicSleeping)
            {
                WakeOnAddress((const void *)&this->uAtomicState);
            }
        }

        inline void Lock() override
        {
            static const AuUInt32 kRef { 1 };

            if (TryLock())
            {
                return;
            }

            while (!TryLockNoSpin())
            {
                AuAtomicAdd(&this->uAtomicSleeping, 1u);
                WaitOnAddress((void *)&this->uAtomicState, &kRef, sizeof(kRef), 0);
                AuAtomicSub(&this->uAtomicSleeping, 1u);
            }
        }

        inline bool LockMS(AuUInt64 qwTimeout) override
        {
            return LockNS(AuMSToNS<AuUInt64>(qwTimeout));
        }

        inline bool LockAbsMS(AuUInt64 qwTimeout) override
        {
            return LockAbsNS(AuMSToNS<AuUInt64>(qwTimeout));
        }

        inline bool LockNS(AuUInt64 qwTimeout) override
        {
            static const AuUInt32 kRef { 1 };

            if (TryLock())
            {
                return true;
            }

            auto qwEndTime = Time::SteadyClockNS() + qwTimeout;

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

                AuAtomicAdd(&this->uAtomicSleeping, 1u);
                bStatus = WaitOnAddressSteady((void *)&this->uAtomicState, &kRef, sizeof(kRef), qwEndTime);
                AuAtomicSub(&this->uAtomicSleeping, 1u);

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

            return true;
        }

        inline bool LockAbsNS(AuUInt64 qwTimeoutAbs) override
        {
            static const AuUInt32 kRef { 1 };
       
            if (TryLock())
            {
                return true;
            }

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

                AuAtomicAdd(&this->uAtomicSleeping, 1u);
                bStatus = WaitOnAddressSteady((void *)&this->uAtomicState, &kRef, sizeof(kRef), qwTimeoutAbs);
                AuAtomicSub(&this->uAtomicSleeping, 1u);

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

            return true;
        }

        volatile AuUInt32 uAtomicState {};
        volatile AuUInt32 uAtomicSleeping {};
    };
}