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

int lock = 0;
int VAR = 0;
int niter = 0;

void counter_worker(int *arg)
{
    int tid = *arg;
    free(arg);

    int step = niter / 4;
    if (step == 0)
        step = 1;

    for (int i = 0; i < niter; i++) {
        futex(&lock, FUTEX_LOCK, 1);

        if (i > 0 && i % step == 0) {
            printf("[T-%d] %d/%d iterations done\n", tid, i, niter);
        }

        int tmp;

        tmp = VAR;
        tmp -= 5;
        VAR = tmp;

        tmp = VAR;
        tmp += 2;
        VAR = tmp;

        tmp = VAR;
        tmp += 4;
        VAR = tmp;

        futex(&lock, FUTEX_UNLOCK, 1);
    }

    futex(&lock, FUTEX_LOCK, 1);
    printf("[T-%d] Finished\n", tid);
    futex(&lock, FUTEX_UNLOCK, 1);

    exit(0);
}

void sleeplock_worker(int *arg)
{
    int tid = *arg;
    free(arg);

    futex(&lock, FUTEX_LOCK, 1);
    printf("[pid %d] T-%d started\n", getpid(), tid);
    futex(&lock, FUTEX_UNLOCK, 1);

    // let all threads start so they contend on the same lock
    pause(2);

    for (int i=0; i<2; i++) {
      futex(&lock, FUTEX_LOCK, 1);
      printf("====== [pid %d] T-%d ACQUIRED LOCK ====== +\n", getpid(), tid);
  
      // hold the lock long enough to force others to block
      pause(8);
  
      printf("====== [pid %d] T-%d RELEASING LOCK ===== -\n", getpid(), tid);
      futex(&lock, FUTEX_UNLOCK, 1);
  
      pause(2);
    }


    futex(&lock, FUTEX_LOCK, 1);
    printf("[pid %d] T-%d exiting\n", getpid(), tid);
    futex(&lock, FUTEX_UNLOCK, 1);

    exit(0);
}

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

    int nthreads = atoi(argv[1]);
    niter = atoi(argv[2]);

    printf("[main] === Test A: sleeplock correctness ===\n");
    printf("[main] initial VAR = %d, nthreads = %d\n", VAR, nthreads);

    for (int i = 0; i < nthreads; i++) {
        int *thread_id = malloc(sizeof(int));
        *thread_id = i;

        create_thread(counter_worker, thread_id);
    }

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

    printf("[main] all counter threads joined\n");
    printf("[main] VAR = %d, expected = %d\n\n", VAR, nthreads * niter);

    printf("[main] === Test B: sleeplock scheduling / blocking behavior ===\n");
    printf("[main] Launching %d threads\n", nthreads);

    set_sched_trace(1);

    for (int i = 0; i < nthreads; i++) {
        int *thread_id = malloc(sizeof(int));
        *thread_id = i;

        create_thread(sleeplock_worker, thread_id);
    }

    for (int i = 0; i < nthreads; i++) {
        join();
    }
    set_sched_trace(0);

    printf("\n[main] all threads joined\n");
    exit(0);
}