#include <addr.h>
#include <printf.h>
#include <madt.h>
#include <timer.h>
#include <cpuid.h>
#include <int.h>
#include <libc.h>
#include <paging.h>
#include <heap.h>
#include <panic.h>

#define LAPIC_ICR_LOW   192
#define LAPIC_ICR_HIGH  196

extern lapic_t lapic;
extern char __load_start_smp_bootloader, __load_stop_smp_bootloader;
extern uint8_t *smp_bootstrap_corecount;
extern uint8_t smp_bootstrap_bsp;
extern struct gdt_descriptor final_gdt_descriptor;
extern struct cpu_info *smp_stackarray_ptr;

uint8_t corecount; //total corecount

struct icr_reg {
  uint8_t vector;
  unsigned int deliv_mode:3;
  unsigned int dest_mode:1;
  unsigned int deliv_status:1;
  unsigned int reserved:1;
  unsigned int level:1; //0 = de-assert
  unsigned int trig_mode:1;
  unsigned int reserved_1:2;
  unsigned int dest_shorthand:2;
  unsigned int reserved_2:12;
}__attribute__((packed));


static inline void write_icr(uint8_t dest, uint32_t message) {
  lapic[LAPIC_ICR_HIGH] = (uint32_t)dest << 24;
  lapic[LAPIC_ICR_LOW] = message;
}

struct gdt_descriptor {
  uint16_t size;
  void *offset;
} __attribute__((packed));

struct cpu_info {
  uint8_t apic_id;
  uint8_t secondary_bsp; //force 8 bits for alignment and consistency purposes
  void *stack;
} __attribute__((packed));


static struct gdt_descriptor gdtr;
struct cores_info cores;

struct apicid_to_coreid_deleteme { //WILL BE DELETED AFTER THREADING EXISTS TODO
  uint8_t apic_id;
  uint8_t core_id;
};

static struct apicid_to_coreid_deleteme *apicid_to_coreid;


uint8_t get_coreid() { //WILL BE DELETED AFTER THREADING EXISTS TODO
  uint32_t ebx, unused_cpuid;
  uint8_t apic_id;
  __get_cpuid(1, &unused_cpuid, &ebx, &unused_cpuid, &unused_cpuid); 
  apic_id = ebx >> 24;
  for(uint8_t core = 0; core < corecount; core++) {
    if(apicid_to_coreid[core].apic_id == apic_id) return apicid_to_coreid[core].core_id;
  }
  return 0;
}

void smp_prepare() {
  uint8_t cores_active = 1;
  uint8_t stack_i = 0, core_i;
  struct icr_reg icr;
  struct cpu_info *stackarray;
  get_coreinfo(&cores);
  if(cores.corecount == 1) {
    asm("sgdt [%0]"::"m"(gdtr));
    gdtr.offset = PHYS_TO_VIRT(gdtr.offset);
    asm("lgdt [%0]"::"m"(gdtr));
    return;
  }
  bzero(&icr, sizeof(icr));

  corecount = cores.corecount;
 
  stackarray = malloc(sizeof(struct cpu_info) * (cores.corecount - 1));
  apicid_to_coreid = malloc(sizeof(struct apicid_to_coreid_deleteme) * (cores.corecount - 1));
  for(core_i = 0; core_i < cores.corecount; core_i++) {

    //WILL BE DELETED AFTER THREADING EXISTS TODO
    apicid_to_coreid[core_i].apic_id = cores.apic_id[core_i];
    apicid_to_coreid[core_i].core_id = core_i;

    if(cores.apic_id[core_i] == cores.bsp) continue;
    stackarray[stack_i].apic_id = cores.apic_id[core_i];
    stackarray[stack_i].stack = palloc(0x1000);
    printf("core %i's stack: %lx\n", core_i, stackarray[stack_i].stack);
    stackarray[stack_i].secondary_bsp = (stack_i)? false : true;

    stack_i++;
  }
  for(stack_i = 0; stack_i < (cores.corecount - 1); stack_i++) {
    printf("%i's stack: %lx\n", stackarray[stack_i].apic_id, stackarray[stack_i].stack);
  }

  memcpy(PHYS_TO_VIRT((void *)0x8000), PHYS_TO_VIRT(&__load_start_smp_bootloader), 
      &__load_stop_smp_bootloader - &__load_start_smp_bootloader);

  //we have to use the virtual address, even though lowmem is still mapped,
  //so GCC doesn't complain about the mcmodel
  *(uint8_t *)PHYS_TO_VIRT(&smp_bootstrap_bsp) = cores.bsp;
  *(uint8_t **)PHYS_TO_VIRT(&smp_bootstrap_corecount) = &cores_active;
  *(struct cpu_info **)PHYS_TO_VIRT(&smp_stackarray_ptr) = stackarray;

  icr.deliv_mode = 0b101;
  icr.dest_shorthand = 0b11;
  icr.level = 1;
  write_icr(0, *(uint32_t *)&icr);
  usleep(10000);

  icr.deliv_mode = 0b110;
  icr.vector = 8;
  write_icr(0, *(uint32_t *)&icr);
  usleep(200);
  if(cores_active != cores.corecount) write_icr(0, *(uint32_t *)&icr);
  usleep(200);
  if(cores_active != cores.corecount) {
    printf("Only %i cores online (expected %i)\n", cores_active, cores.corecount); 
    PANIC(KERNEL_PANIC_SMP_FAILED); //will replace this with warning once we push public
  }
  printf("All detected %i cores have booted\n", cores_active);
  asm("lgdt [%0]\n"::"r"(PHYS_TO_VIRT(&final_gdt_descriptor))); //note that segment registers are not reloaded
}