#include <amd64/apic.h>
#include <amd64/init.h>
#include <amd64/mm.h>
#include <amd64/msr-index.h>
#include <amd64/msr.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/liballoc.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/smp.h>
#include <sys/syscall.h>

/// Cpu ID counter
static uint32_t cpu_counter = 0;
/// Lock for \ref cpu_counter
static spin_lock_t cpu_counter_lock = SPIN_LOCK_INIT;
/// The CPUs
static struct cpu cpus[CPUS_MAX];

static bool thiscpu_init = false;

void amd64_thiscpu_set_init (void) { thiscpu_init = true; }

/// Allocate a CPU structure
struct cpu* cpu_make (void) {
  spin_lock (&cpu_counter_lock);
  int id = cpu_counter++;
  spin_unlock (&cpu_counter_lock);

  struct cpu* cpu = &cpus[id];

  memset (cpu, 0, sizeof (*cpu));
  cpu->lock = SPIN_LOCK_INIT;
  cpu->id = id;

  amd64_wrmsr (MSR_SHADOW_GS_BASE, (uint64_t)cpu);

  return cpu;
}

struct cpu* cpu_get (void) {
  if (!thiscpu_init)
    return NULL;

  return (struct cpu*)amd64_rdmsr (MSR_SHADOW_GS_BASE);
}

/// Bootstrap code for non-BSP CPUs
static void amd64_smp_bootstrap (struct limine_mp_info* mp_info) {
  amd64_load_kernel_cr3 ();

  struct cpu* cpu = cpu_make ();

  amd64_init (cpu, true); /* gdt + idt */
  syscall_init ();

  thiscpu->lapic_ticks = amd64_lapic_init (2500);
  amd64_lapic_tick (thiscpu->lapic_ticks);

  DEBUG ("CPU %u is online!\n", thiscpu->id);

  __asm__ volatile ("sti");

  for (;;)
    ;
}

/// Initialize SMP subsystem for AMD64. Start AP CPUs
void smp_init (void) {
  thiscpu->lapic_ticks = amd64_lapic_init (2500);

  struct limine_mp_response* mp = limine_mp_request.response;

  for (size_t i = 0; i < mp->cpu_count; i++) {
    if (mp->cpus[i]->lapic_id != thiscpu->id) {
      DEBUG ("Trying CPU %u\n", mp->cpus[i]->lapic_id);
      mp->cpus[i]->goto_address = &amd64_smp_bootstrap;
    }
  }
}