#include <stdatomic.h>
#include <stdint.h>
#include <stdbool.h>
#include "spinlock.h"
#include "hal/hal.h"
#include "kprintf.h"

#define SPINLOCK_HINT() asm volatile("pause")

struct {
  uint64_t irq_flags;
  atomic_int irq_nest;
} IRQ_CTX = {0};

void spinlock_init(SpinLock *sl) {
  atomic_clear_flag_explicit(&sl->lock, memory_order_relaxed);
}

uint64_t irqsave(void) {
  uint64_t flags;
  asm volatile("pushfq; cli; popq %0" : "=r"(flags) :: "memory", "cc");
  return flags;
}

void irqrestore(uint64_t flags) {
  if (flags & (1ULL << 9)) {
    asm volatile("sti" ::: "memory", "cc");
  }
}

void irqsave_nested(void) {
  int prev = atomic_fetch_add_explicit(&IRQ_CTX.irq_nest, 1, memory_order_acq_rel);
  if (prev == 0) {
    IRQ_CTX.irq_flags = irqsave();
  }
}

void irqrestore_nested(void) {
  int prev = atomic_fetch_sub_explicit(&IRQ_CTX.irq_nest, 1, memory_order_acq_rel);
  if (prev == 1) {
    irqrestore(IRQ_CTX.irq_flags);
  }
}

void spinlock_acquire(SpinLock *sl) {
  irqsave_nested();

  while (atomic_test_and_set_explicit(&sl->lock, memory_order_acquire)) {
    SPINLOCK_HINT();
  }
}

void spinlock_release(SpinLock *sl) {
  atomic_clear_flag_explicit(&sl->lock, memory_order_release);

  irqrestore_nested();
}