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

    File: AuRWLock.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuRWLock.hpp"

namespace Aurora::Threading::Primitives
{
    template<bool isread>   
    void RWLockAccessView<isread>::Unlock()
    {
        if constexpr (isread)
        {
            this->parent_.UnlockRead();
        }
        else
        {
            this->parent_.UnlockWrite();
        }
    }

    template<bool isread>
    bool RWLockAccessView<isread>::Lock(AuUInt64 timeout)
    { 
        if constexpr (isread)
        {
            return this->parent_.LockRead(timeout);
        }
        else
        {
            return this->parent_.LockWrite(timeout);
        }
    }

    template<bool isread>
    bool RWLockAccessView<isread>::TryLock()
    {
        if constexpr (isread)
        {
            return this->parent_.TryLockRead();
        }
        else
        {
            return this->parent_.TryLockWrite();
        }
    }

    RWLockImpl::RWLockImpl() : read_(*this), write_(*this)
    {

    }

    RWLockImpl::~RWLockImpl()
    {
        this->mutex_.reset();
        this->condition_.reset();
    }

    bool RWLockImpl::Init()
    {
        this->mutex_ = ConditionMutexUnique();
        if (!this->mutex_)
        {
            return false;
        }

        this->condition_ = ConditionVariableUnique(AuUnsafeRaiiToShared(mutex_));
        if (!this->condition_)
        {
            return false;
        }
        return true;
    }
        
    bool RWLockImpl::LockRead(AuUInt64 timeout)
    {
    #if 0
        AU_LOCK_GUARD(mutex_);

        if (this->state_ == -1 && this->reentrantWriteLockHandle_ == AuThreads::GetThreadId())
        {
            return true;
        }

        while (this->state_ < 0 /* || this->writersPending_*/)
        {
            if (!this->condition_->WaitForSignal(timeout))
            {
                return false;
            }

            if (this->writersPending_)
            {
                // Meh, let's just DoS the cpu with the readers until we find the writer for low unlock-from-final-read latency
                // The writer should be prio, and it's already having a terrible day by stalling. Let's not play bounce the signal through contexts (we were likely a signal, not a broadcast)
                this->condition_->Broadcast();
                continue;
            }
        }

        this->state_++;
    #else

        if (this->state_ == -1 && this->reentrantWriteLockHandle_ == AuThreads::GetThreadId())
        {
            return true;
        }

        AuInt32 iCurState {};
        do
        {
            iCurState = this->state_;

            if (iCurState < 0)
            {
                AU_LOCK_GUARD(this->mutex_);

                iCurState = this->state_;
                if (iCurState < 0)
                {
                    if (!this->condition_->WaitForSignal(timeout))
                    {
                        return false;
                    }

                    if (this->writersPending_)
                    {
                        this->condition_->Broadcast();
                        continue;
                    }
                }
            }
        }
        while (iCurState == -1 ||
               AuAtomicCompareExchange(&this->state_, iCurState + 1, iCurState) != iCurState);

    #endif
        return true;
    }

    bool RWLockImpl::LockWrite(AuUInt64 timeout)
    {
        if (AuAtomicCompareExchange(&this->state_, -1, 0) == 0)
        {
            this->reentrantWriteLockHandle_ = AuThreads::GetThreadId();
            return true;
        }

        AU_LOCK_GUARD(this->mutex_);
        this->writersPending_++;

        while (this->state_ != 0)
        {
            if (!this->condition_->WaitForSignal(timeout))
            {
                this->writersPending_--;
                return false;
            }
        }

        this->reentrantWriteLockHandle_ = AuThreads::GetThreadId();
        this->writersPending_--;
        this->state_ = -1;
        return true;
    }
    
    bool RWLockImpl::TryLockRead()
    {
        auto iCurState = this->state_;

        if (iCurState == -1)
        {
            return this->reentrantWriteLockHandle_ == AuThreads::GetThreadId();
        }

        return AuAtomicCompareExchange(&this->state_, iCurState + 1, iCurState) == iCurState;
    }

    bool RWLockImpl::TryLockWrite()
    {
        AU_LOCK_GUARD(this->mutex_);
    
        if (this->state_ > 0)
        {
            return false;
        }

        this->reentrantWriteLockHandle_ = AuThreads::GetThreadId();
        this->state_ = -1;
        return true;
    }
    
    void RWLockImpl::UnlockRead()
    {
        AU_LOCK_GUARD(this->mutex_);

        if (this->state_ == -1)
        {
            SysAssertDbg(this->reentrantWriteLockHandle_ == AuThreads::GetThreadId());
            return;
        }

        auto val = AuAtomicSub(&this->state_, 1);
        
        if ((val == 1) && (this->bElevaterPending_))
        {
            this->condition_->Signal();
        }

        if (val == 0)
        {
            this->condition_->Signal();
        }
    }

    void RWLockImpl::UnlockWrite()
    {
        AU_LOCK_GUARD(this->mutex_);
        this->state_ = 0;
        this->condition_->Broadcast();
        this->reentrantWriteLockHandle_ = 0;
    }
    
    bool RWLockImpl::UpgradeReadToWrite(AuUInt64 timeout)
    {
        AU_LOCK_GUARD(this->mutex_);
    
        while (this->state_ != 1)
        {
            this->bElevaterPending_ = true;

            if (!this->condition_->WaitForSignal(timeout))
            {
                return false;
            }
        }
        
        this->bElevaterPending_ = false;
        this->reentrantWriteLockHandle_ = AuThreads::GetThreadId();
        this->state_ = -1;
        return true;
    }

    bool RWLockImpl::DowngradeWriteToRead()
    {
        AU_LOCK_GUARD(this->mutex_);

        if (this->state_ != -1)
        {
            return false;
        }

        this->state_ = 1;
        this->condition_->Broadcast();
        return true;
    }

    IWaitable *RWLockImpl::AsReadable()
    {
        return &this->read_;
    }

    IWaitable *RWLockImpl::AsWritable()
    {
        return &this->write_;
    }
    
    AUKN_SYM IRWLock *RWLockNew()
    {
        auto pRwLock = _new RWLockImpl();
        if (!pRwLock)
        {
            return nullptr;
        }

        if (!pRwLock->Init())
        {
            delete pRwLock;
            return nullptr;
        }

        return pRwLock;
    }

    AUKN_SYM void RWLockRelease(IRWLock *pRwLock)
    {
        AuSafeDelete<RWLockImpl *>(pRwLock);
    }
}