#include <limine.h>
#include <logging.h>
#include <kernel/pmm.h>
#include <arch/cpu.h>
#include <arch/gdt.h>
#include <arch/idt.h>
#include <kernel/kernel_args.h>
#include <arch/smp.h>
#include <kernel/alloc.h>
#include <kernel/init.h>
#include <arch/context.h>
#include <arch/prof.h>

DEFINE_KERNEL_ARGUMENT(nosmp, bool);

static volatile struct limine_mp_request smprequest = {
	.id = LIMINE_MP_REQUEST_ID
};

// set to 1 because of the bsp
size_t arch_smp_cpusawake = 1;

// for panic. if the nosmp argument is given to the kernel, this is what the APs will jump to
static void cpuwakeuphalt(struct limine_mp_info *info) {
	asm("cli");
	for (;;) CPU_HALT();
}

static long next_id = 1;

static void cpuwakeup(struct limine_mp_info *info) {
	cpu_set((cpu_t *)info->extra_argument);

	current_cpu()->internal_id = __atomic_fetch_add(&next_id, 1, __ATOMIC_SEQ_CST);

	arch_gdt_reload();
	arch_idt_reload();
	interrupt_register(0xfd, (void *)cpuwakeuphalt, NULL, IPL_IGNORE);
	dpc_init();
	arch_extracontext_detect();
	arch_mmu_apswitch();
	vmm_apinit();

	arch_apic_initap();

	timekeeper_early_init();

	arch_cpu_init();
	arch_apic_timerinit();
	
	timekeeper_init();

#ifdef ENABLE_PROFILING
	arch_profiling_init();
#endif

	__atomic_add_fetch(&arch_smp_cpusawake, 1, __ATOMIC_SEQ_CST);
	sched_ap_entry();
}

void arch_smp_send_ipi(cpu_t *targcpu, isr_t *isr, int target, bool nmi) {
	arch_apic_sendipi(targcpu ? targcpu->hardware_id : 0, INTERRUPT_IDTOVECTOR(isr->id), target, nmi ? APIC_MODE_NMI : 0, 0);
}

void arch_smp_haltallothers(void) {
	arch_smp_send_ipi(NULL, &current_cpu()->isr[0xfd], ARCH_SMP_IPI_OTHERCPUS, true);
}

size_t arch_smp_get_cpu_count(void) {
	if (GET_KERNEL_ARGUMENT(nosmp, bool))
		return 1;

	struct limine_mp_response *response = smprequest.response;
	if (response == NULL)
		return 1;

	return response->cpu_count;
}

cpu_t **smp_cpus;

cpu_t *arch_smp_get_cpu_by_internal_id(long id) {
	__assert(smp_cpus && id < arch_smp_get_cpu_count());
	return smp_cpus[id];
}

void arch_smp_wakeup(void) {
	struct limine_mp_response *response = smprequest.response;
	if (response == NULL) {
		// TODO try to detect the other processors manually
		printf("smp: limine smp request response not present\n");
		return;
	}

	printf("smp: %d processor%s\n", response->cpu_count, response->cpu_count > 1 ? "s" : "");

	// use physical pages so the other cpus have it on the hhdm
	size_t cpu_size = ROUND_UP(sizeof(cpu_t) * response->cpu_count, PAGE_SIZE);
	cpu_t *cpus = pmm_alloc(cpu_size / PAGE_SIZE, PMM_SECTION_DEFAULT);
	__assert(cpus);
	cpus = MAKE_HHDM(cpus);
	memset(cpus, 0, cpu_size);

	smp_cpus = alloc(response->cpu_count * sizeof(cpu_t *));
	__assert(smp_cpus);

	void (*wakeupfn)(struct limine_mp_info *) = GET_KERNEL_ARGUMENT(nosmp, bool) ? cpuwakeuphalt : cpuwakeup;

	// make the other processors jump to cpuwakeup()
	for (int i = 0; i < response->cpu_count; ++i) {
		// skip the bootstrap processor
		if (response->cpus[i]->lapic_id == response->bsp_lapic_id) {
			smp_cpus[i] = get_bsp();
			continue;
		}

		smp_cpus[i] = &cpus[i];
		response->cpus[i]->extra_argument = (uint64_t)&cpus[i];

		__atomic_store_n(&response->cpus[i]->goto_address, wakeupfn, __ATOMIC_SEQ_CST);
	}

	// wait for other cpus to boot up
	if (wakeupfn == cpuwakeup)
		while (__atomic_load_n(&arch_smp_cpusawake, __ATOMIC_SEQ_CST) != response->cpu_count) asm("pause");

	printf("smp: awoke other processors\n");
}

INIT_ROUTINE_DEFINE(smp, INIT_ROUTINE_FLAGS_NONE, arch_smp_wakeup, scheduler);