#include "kernel/types.h"
#include "kernel/stat.h"
#include "user/user.h"

static int lock = 0;

void worker(int *arg)
{
    futex(&lock, FUTEX_LOCK, 1);

    /*
     * In this test, main controls who gets woken up.
     * Workers only release after they acquire the lock.
     *
     * Main thread is making sure that there are no deadlocks.
     */

    futex(&lock, FUTEX_UNLOCK, 0);
    exit(0);
}

static void
spawn_workers(int nthreads)
{
    for (int i = 0; i < nthreads; i++) {
        create_thread(worker, 0);
    }
}

static void
join_workers(int nthreads)
{
    for (int i = 0; i < nthreads; i++) {
        join();
    }
}

int main(int argc, char *argv[])
{
    if (argc != 2) {
        printf("Usage: %s <num_threads>\n", argv[0]);
        exit(1);
    }

    int nthreads = atoi(argv[1]);

    /* Test A */
    printf("\n=== Test A: wake all waiters ===\n");

    futex(&lock, FUTEX_LOCK, 1);
    printf("[pid %d main] launching %d threads\n", getpid(), nthreads);

    set_sched_trace(1);
    spawn_workers(nthreads);

    // Give workers time to block on the lock.
    pause(5);

    printf("[pid %d main] waking all threads\n", getpid());
    futex(&lock, FUTEX_UNLOCK, 1);

    printf("[pid %d main] all threads woken\n", getpid());

    /*  All worker threads must have been scheduled once before this.

       The following loop is here to avoid deadlock.
       In case any workers tried to acquire the lock twice
       and go to sleep again, main thread needs to wake it
       up.
    */
    for (int i = 0; i < nthreads - 1; i++) {
        futex(&lock, FUTEX_LOCK, 0);
        futex(&lock, FUTEX_UNLOCK, 1);
    }

    join_workers(nthreads);
    set_sched_trace(0);

    /* Test B */
    printf("\n=== Test B: wake one waiter at a time ===\n");

    futex(&lock, FUTEX_LOCK, 1);
    printf("[pid %d main] launching %d threads again\n", getpid(), nthreads);

    set_sched_trace(1);
    spawn_workers(nthreads);

    // Give workers time to block on the lock.
    pause(5);

    for (int i = 0; i < nthreads; i++) {
        printf("[pid %d main] waking one thread\n", getpid());
        futex(&lock, FUTEX_UNLOCK, 2);

        // let this thread end execution.
        pause(5);

        // Re-acquire so main remains the one controlling the next wake.
        if (i != nthreads - 1)
            futex(&lock, FUTEX_LOCK, 1);
    }

    join_workers(nthreads);
    set_sched_trace(0);

    printf("\n[pid %d main] all threads joined\n", getpid());
    exit(0);
}
