#include <assert.h>
#include <atomic>
#include <cstdint>
#include <iostream>
#include <memory>
#include <shared_mutex>
#include <thread>
#include <vector>

#if 0
constexpr const int kNumThreads = 5;

struct ThreadTestData {
    std::shared_mutex sharedMutex = {};
    std::atomic<int32_t> readCounter;
};

int DoStuff(ThreadTestData* data) {
    // Acquire reader lock
    data->sharedMutex.lock_shared();

    // wait until all read threads have acquired their shared lock
    data->readCounter.fetch_add(1);
    while (data->readCounter.load() != kNumThreads) {
        std::this_thread::yield();
    }

    // Release reader lock
    data->sharedMutex.unlock_shared();

    return 0;
}

int main() {
    int count = 0;
    while (true) {
        ThreadTestData data = {};

        // Acquire write lock
        data.sharedMutex.lock();

        // Create N threads
        std::vector<std::unique_ptr<std::thread>> readerThreads;
        readerThreads.reserve(kNumThreads);
        for (int i = 0; i < kNumThreads; ++i) {
            readerThreads.emplace_back(std::make_unique<std::thread>(DoStuff, &data));
        }

        // Release write lock
        data.sharedMutex.unlock();

        // Wait for all readers to succeed
        for (auto& thread : readerThreads) {
            thread->join();
        }

        // Cleanup
        readerThreads.clear();

        // Spew so we can tell when it's deadlocked
        count += 1;
        std::cout << count << std::endl;
    }

    return 0;
}
#else
#include <Windows.h>
#include <winternl.h>

using NtYieldExecutionFunc = uint32_t(*)(void);
NtYieldExecutionFunc ntYieldExecution = nullptr;

struct ThreadTestData
{
    HANDLE hEvent;
    SRWLOCK SRWLock = {};
    LONG numThreads = 1;
    LONG readCounter = 0;

    void EndThread()
    {
        if (!InterlockedDecrementNoFence(&numThreads))
        {
            if (!SetEvent(hEvent)) __debugbreak();
        }
    }

    void DoStuff()
    {
        AcquireSRWLockShared(&SRWLock);

        ULONG n = 0, m = 0, k = InterlockedDecrementNoFence(&readCounter);

        while (readCounter)
        {
            switch (ntYieldExecution())
            {
            case /*STATUS_SUCCESS*/0:
                n++;
                break;
            case /*STATUS_NO_YIELD_PERFORMED*/0x40000024L:
                m++;
                break;
            default:
                __debugbreak();
            }
        }

        ReleaseSRWLockShared(&SRWLock);

        //std::cout << k << ">" << n << "/" << m << std::endl;

        EndThread();
    }

    static ULONG WINAPI _S_DoStuff(PVOID data)
    {
        reinterpret_cast<ThreadTestData*>(data)->DoStuff();
        return 0;
    }

    void Test(ULONG n)
    {
        if (hEvent = CreateEventW(0, 0, 0, 0))
        {
            AcquireSRWLockExclusive(&SRWLock);

            do
            {
                numThreads++;
                readCounter++;

                if (HANDLE hThread = CreateThread(0, 0, _S_DoStuff, this, 0, 0))
                {
                    CloseHandle(hThread);
                }
                else
                {
                    readCounter--;
                    numThreads--;
                }

            } while (--n);

            ReleaseSRWLockExclusive(&SRWLock);

            EndThread();

            if (WAIT_OBJECT_0 != WaitForSingleObject(hEvent, INFINITE))
            {
                __debugbreak();
            }

            CloseHandle(hEvent);
        }
    }
};

int main()
{
    auto ntdll = LoadLibraryA("ntdll.dll");
    assert(ntdll);
    ntYieldExecution = (NtYieldExecutionFunc)GetProcAddress(ntdll, "NtYieldExecution");

    int count = 0;
    while (true) {
        ThreadTestData data;
        data.Test(8);

        count += 1;
        std::cout << "Count: " << count << std::endl;
    }

    return 0;
}

#endif