/***
    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 TryLock() override
        {
            if (AuAtomicTestAndSet(&this->uAtomicState, 0u))
            {
                return true;
            }

        #if defined(AURORA_ARCH_X86) || defined(AURORA_ARCH_X64)
            for (AU_ITERATE_N(i, AuUInt(1u) << AuUInt(Threading::GetSpinCountTimeout())))
            {
                _mm_pause();

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

            return false;
        }

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

        inline bool HasLockImplementation() override
        {
            return true;
        }

        inline void Unlock() override
        {
            this->uAtomicState = 0;

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

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

            while (!TryLock())
            {
                AuAtomicAdd(&this->uAtomicSleeping, 1u);
                WaitOnAddress(&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 (AuAtomicTestAndSet(&this->uAtomicState, 0u))
            {
                return true;
            }

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

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

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

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

            return true;
        }

        inline bool LockAbsNS(AuUInt64 qwTimeoutAbs) override
        {
            static const AuUInt32 kRef { 1 };
            if (AuAtomicTestAndSet(&this->uAtomicState, 0u))
            {
                return true;
            }

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

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

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

            return true;
        }

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