//
// Created by Aaron Gill-Braun on 2023-12-22.
//

#include <kernel/lock.h>
#include <kernel/mm.h>
#include <kernel/proc.h>
#include <kernel/tqueue.h>
#include <kernel/panic.h>

#define ASSERT(x) kassert(x)

#define MAX_CLAIMS 8

struct lock_class *lock_classes[NUM_LOCK_CLASSES];

/*
 * A lock claim is a record of a lock held by an owner.
 */
struct lock_claim {
  struct lock_object *lock; // the owned lock
  uintptr_t how;            // how the lock was acquired
  const char *file;         // file where the lock was acquired
  int line;                 // line where the lock was acquired
};

/*
 * A list of lock_claims.
 * It is used to hold a number of lock claims and avoid allocations in the locking
 * path which results in a new claim being written. We only need to allocate when
 * the per-node list is full.
 */
struct lock_claim_list {
  struct lock_claim claims[MAX_CLAIMS];// list of claims
  int nclaims;                  // number of claims
  struct lock_claim_list *next; // next list
};

// MARK: lock claims

static void percpu_early_init_claim_list() {
  // this per-cpu lock lock_claim_list tracks spin lock claims
  struct lock_claim_list *list = lock_claim_list_alloc();
  PERCPU_AREA->spin_claims = list;
}
PERCPU_EARLY_INIT(percpu_early_init_claim_list);


static pool_t *lock_claims_pool;

static void lock_claims_static_init() {
  lock_claims_pool = pool_create("lock_claims", pool_sizes(sizeof(struct lock_claim_list)), 0);
  if (!lock_claims_pool) {
    panic("failed to create lock_claims_pool");
  }
  pool_preload_cache(lock_claims_pool, sizeof(struct lock_claim_list), 8);
}
STATIC_INIT(lock_claims_static_init);


struct lock_claim_list *lock_claim_list_alloc() {
  struct lock_claim_list *list;
  if (__expect_false(lock_claims_pool == NULL)) {
    list = kmallocz(sizeof(struct lock_claim_list));
  } else {
    list = pool_alloc(lock_claims_pool, sizeof(struct lock_claim_list));
  }
  list->nclaims = 0;
  list->next = NULL;
  return list;
}

void lock_claim_list_free(struct lock_claim_list **listp) {
  struct lock_claim_list *list = *listp;
  while (list) {
    struct lock_claim_list *next = list->next;
    if (lock_claims_pool == NULL || kheap_is_valid_ptr(list)) {
      kfree(list);
    } else {
      pool_free(lock_claims_pool, list);
    }
    list = next;
  }
  *listp = NULL;
}

void lock_claim_list_add(struct lock_claim_list *list, struct lock_object *lock, uintptr_t how, const char *file, int line) {
  // traverse to the last node that has space
  while (list->nclaims == MAX_CLAIMS) {
    if (list->next == NULL) {
      // allocate a new node
      list->next = lock_claim_list_alloc();
    }
    list = list->next;
  }

  struct lock_claim *claim = &list->claims[list->nclaims++];
  claim->lock = lock;
  claim->how = how;
  claim->file = file;
  claim->line = line;
}

void lock_claim_list_remove(struct lock_claim_list *list, struct lock_object *lock) {
  // search through all nodes in the linked list
  while (list != NULL) {
    // scan in reverse order to find the most recent claim in this node
    for (int i = list->nclaims - 1; i >= 0; i--) {
      if (list->claims[i].lock == lock) {
        // found the lock, remove it
        list->claims[i].lock = NULL;
        list->claims[i].how = 0;
        list->claims[i].file = NULL;
        list->claims[i].line = 0;
        list->nclaims--;
        
        // shift the rest of the claims down
        for (int j = i; j < list->nclaims; j++) {
          list->claims[j] = list->claims[j + 1];
        }
        return;
      }
    }
    // move to the next node
    list = list->next;
  }

  // if we get here then the lock was not found
  panic("lock_claim_list_remove() on unowned lock");
}

// MARK: spin delay

int spin_delay_wait(struct spin_delay *delay) {
  if (delay->waits >= delay->max_waits) {
    return 0;
  }

  register uint64_t count asm ("r15") = delay->delay_count;
  while (count--) {
    cpu_pause();
  }

  delay->waits++;
  return 1;
}
