#include <arch/acpi.h>
#include <arch/acpi/local_apic.hpp>
#include <socks/printk.h>
#include <socks/pmap.h>
#include <socks/vm.h>
#include <socks/cpu.h>

using namespace arch::acpi;

extern "C" uint8_t acpi_bsp_lapic_id(void);
extern "C" char ap_trampoline[];

volatile struct vm_page *__ap_stack_page = NULL;
volatile uintptr_t __ap_stack_top = 0;
volatile uint8_t __this_ap_ok = 0;
volatile uint8_t __all_ap_ok = 0;
volatile unsigned int __this_ap_id = 0;

extern "C" void ap_trampoline_exit(void)
{
	ml_cpu_block *this_cpu = (ml_cpu_block *)kmalloc(sizeof *this_cpu, VM_NORMAL);
	ml_cpu_block_init(this_cpu);
	ml_cpu_block_use(this_cpu);

	this_cpu->c_cpu_id = __this_ap_id;

	struct cpu_data *self = get_this_cpu();
	self->c_flags = CPU_ONLINE;
	self->c_id = __this_ap_id;
	this_cpu->c_data = self;

	struct thread *this_thread = create_idle_thread();
	this_thread->tr_id = __this_ap_id;
	this_thread->tr_kstack = (struct vm_page *)__ap_stack_page;
	self->c_rq.rq_idle = self->c_rq.rq_cur = this_thread;

	put_cpu(self);

	ap_apic_init();

	__this_ap_ok = 1;
	irq_enable();

	idle();
}

static int init_ap(struct acpi_madt_record *rec, local_apic& lapic, uint8_t bsp_id)
{
	struct lapic_record *ap_lapic = (struct lapic_record *)(rec + 1);
	if (!(ap_lapic->l_flags & 0x1) && !(ap_lapic->l_flags & 0x2)) {
		return -1;
	}

	if (ap_lapic->l_apic_id == bsp_id) {
		//printk("acpi: core %u online [BSP]", ap_lapic->l_apic_id);
		return 0;
	}

	__this_ap_ok = 0;
	__ap_stack_page = vm_page_alloc(VM_PAGE_4K, VM_NORMAL);
	__ap_stack_top = (uintptr_t)vm_page_get_vaddr((struct vm_page *)__ap_stack_page) + VM_PAGE_SIZE;
	__this_ap_id = ap_lapic->l_apic_id;

	lapic.send_ipi(ap_lapic->l_apic_id, 0xC500);
	lapic.send_ipi(ap_lapic->l_apic_id, 0x8500);
	milli_sleep(10);                                                                                                                 // wait 10 msec
	lapic.send_ipi(ap_lapic->l_apic_id, 0x4600 | (AP_TRAMPOLINE_PADDR >> VM_PAGE_SHIFT));

	do {
		ml_cpu_relax();
	} while (__this_ap_ok == 0);

	//printk("acpi: core %u online [AP]", ap_lapic->l_apic_id);
	return 0;
}

kern_status_t bring_all_ap_online(void)
{
	struct acpi_madt *madt = (struct acpi_madt *)acpi_find_sdt(ACPI_SIG_MADT);
	if (!madt) {
		return KERN_UNSUPPORTED;
	}

	uint8_t bsp_id = acpi_bsp_lapic_id();

	void *ap_trampoline_dest = vm_phys_to_virt(AP_TRAMPOLINE_PADDR);
	memcpy(ap_trampoline_dest, ap_trampoline, VM_PAGE_SIZE);

	pmap_t kernel_pmap = get_kernel_pmap();
	uint32_t *pmap_ptr = (uint32_t *)vm_phys_to_virt(0x8ff8);
	*pmap_ptr = kernel_pmap;

	pmap_add(kernel_pmap, (void *)0x8000, 8, (enum vm_prot)(VM_PROT_READ | VM_PROT_EXEC), PMAP_NORMAL);

	unsigned char *p = (unsigned char *)madt + sizeof *madt;
	unsigned char *madt_end = (unsigned char *)madt + madt->m_base.s_length;

	unsigned int nr_processors = 0;

	while (p < madt_end) {
		struct acpi_madt_record *rec = (struct acpi_madt_record *)p;

		switch (rec->r_type) {
		case ACPI_MADT_LAPIC:
			if (init_ap(rec, local_apic::get(), bsp_id) == 0) {
				nr_processors++;
			}
			break;
		default:
			break;
		}

		p += rec->r_length;
	}

	printk("acpi: found %u logical cores", nr_processors);
	return KERN_OK;
}

static void no_smp_config(void)
{
	cpu_set_available(0);
	cpu_set_online(0);

	ml_cpu_block *self = ml_this_cpu();
	self->c_cpu_id = 0;
}

kern_status_t acpi_scan_cpu_topology(void)
{
	struct acpi_madt *madt = (struct acpi_madt *)acpi_find_sdt(ACPI_SIG_MADT);
	if (!madt) {
		no_smp_config();
		return KERN_UNSUPPORTED;
	}

	uint8_t bsp_id = acpi_bsp_lapic_id();

	//void *bsp_lapic = find_lapic(madt);

	unsigned char *p = (unsigned char *)madt + sizeof *madt;
	unsigned char *madt_end = (unsigned char *)madt + madt->m_base.s_length;

	unsigned int nr_processors = 0;

	while (p < madt_end) {
		struct acpi_madt_record *rec = (struct acpi_madt_record *)p;

		if (rec->r_type != ACPI_MADT_LAPIC) {
			p += rec->r_length;
			continue;
		}

		struct lapic_record *lapic = (struct lapic_record *)(rec + 1);

		if (!(lapic->l_flags & 0x1) && !(lapic->l_flags & 0x2)) {
			/* processor cannot be used */
			continue;
		}

		cpu_set_available(lapic->l_apic_id);
		nr_processors++;

		p += rec->r_length;
	}

	cpu_set_online(bsp_id);

	/* store the BSP ID in the current cpu block */
	ml_cpu_block *self = ml_this_cpu();
	self->c_cpu_id = bsp_id;

	printk("acpi: found %u logical cores", nr_processors);
	return KERN_OK;
}

kern_status_t smp_init(void)
{
	return bring_all_ap_online();
}