#include "lib/Mutex.h"
#include "Log.h"
#include "arch/CPU.h"
#include "thread/Scheduler.h"

void Mutex::lock()
{
    auto* current = Scheduler::current();
    const pid_t desired = current->id;
    check(desired > 0); // Why the hell would the idle thread be touching a mutex?

    while (true)
    {
        // Make sure only one thread is touching the mutex at the same time.
        m_spinlock.lock();

        pid_t expected = 0;
        if (!m_thread.compare_exchange_strong(expected, desired))
        {
            if (expected == desired)
            {
                kerrorln("DEADLOCK! KMutex::lock() recursively called by the same thread (%d)", current->id);
                fail("Mutex deadlock detected");
            }

            bool ok = m_blocked_threads.try_push(current);
            m_spinlock.unlock();
            if (!ok) kernel_sleep(10);
            else
                kernel_wait_for_event();
        }
        else
        {
            m_spinlock.unlock();
            break;
        }
    }
};

void Mutex::unlock()
{
    auto* current = Scheduler::current();
    pid_t expected = current->id;
    check(expected > 0); // Why the hell would the idle thread be touching a mutex?

    m_spinlock.lock();

    if (!m_thread.compare_exchange_strong(expected, 0))
    {
        kerrorln("KMutex::unlock() called on a lock already locked by another thread (%d, current is %d)", expected,
                 current->id);
        fail("Mutex unlock by different thread");
    }

    Thread* blocked;
    if (m_blocked_threads.try_pop(blocked))
    {
        while (blocked->state == ThreadState::Runnable)
        {
            // Rarely, we could switch to here between m_spinlock.unlock() and kernel_wait_for_event() above.
            // Let the thread go to sleep before waking it up again.
            kernel_yield();
        }
        blocked->wake_up();
    }

    m_spinlock.unlock();
}

bool Mutex::try_lock()
{
    auto* current = Scheduler::current();
    const pid_t desired = current->id;
    check(desired > 0); // Why the hell would the idle thread be touching a mutex?

    // Make sure only one thread is touching the mutex at the same time.
    m_spinlock.lock();

    pid_t expected = 0;
    bool ok = m_thread.compare_exchange_strong(expected, desired);

    if (expected == desired)
    {
        kwarnln("Deadlock avoided! KMutex::try_lock() failed because it was already locked by the same thread "
                "(%d), this is not supposed to happen",
                current->id);
        CPU::print_stack_trace();
    }

    m_spinlock.unlock();

    return ok;
}