// SPDX-License-Identifier: MPL-2.0
#include "cpu/ap.h"
#include "devs/apic/lapic.h"
#include "devs/input/kbd.h"
#include "memory/memory.h"
#include "kernel/log.h"
#include "std/stdlib.h"
sh_ap_AP_BOOTSTRAP *ap_bootstrap_struct=(sh_ap_AP_BOOTSTRAP*)SH_AP_BOOTSTRAP_PA;
sh_uint8 *is_ap_started=SH_NULLPTR;
static sh_ap_CPU_STRUCT bsp_cpu_struct;
static sh_ap_PER_CPU bsp_per_cpu;
SH_STATUS sh_ap_load_gdt_32() {
  if (sh_page_is_allocated((sh_uint8*)sh_page_get_physical_bitmap_ptr(),SH_AP_GDT_32_PA)) return SH_STATUS_OUT_OF_MEMORY;
  SH_STATUS status=sh_memory_identity_map(SH_AP_GDT_32_PA,1,SH_PAGE_RW | SH_PAGE_PRESENT | SH_PAGE_NX);
  if (sh_status_error(status)) return status;
  sh_gdt_GDT_32 *gdt_32=(sh_gdt_GDT_32*)(sh_uint8*)SH_AP_GDT_32_PA;
  status=sh_gdt_fill_gdt_32(gdt_32);
  if (sh_status_error(status)) return status;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_ap_prepare_for_smp_launch(sh_lapic_DEVICE **lapic_dev_array,sh_uint64 max_lapic_id,sh_uint64 expected_lapic_found,sh_conf_BOOT_CONFIG *boot_config,sh_idt_IDT *idt) {
  if (lapic_dev_array==SH_NULLPTR || expected_lapic_found==0 || max_lapic_id>255) return SH_STATUS_INVALID_PARAMETER;
  SH_STATUS status=sh_memory_identity_map(SH_AP_BOOTSTRAP_PA,1,SH_PAGE_PRESENT | SH_PAGE_RW | SH_PAGE_NX);
  if (sh_status_error(status)) return status;
  sh_ap_AP_BOOTSTRAP *ap_bootstrap=(sh_ap_AP_BOOTSTRAP*)(sh_uint8*)SH_AP_BOOTSTRAP_PA;
  ap_bootstrap->c_entry_point_va=(sh_uint64)&sh_ap_entry_point;
  ap_bootstrap->page_table_pa=boot_config->page_table_pool_pa;
  sh_int16 current_lapic_id=sh_lapic_get_current_core();
  if (current_lapic_id<0) return SH_STATUS_INVALID_STATE;
  sh_lapic_DEVICE *lapic_device_default=SH_NULLPTR;
  sh_uint64 lapic_found=0;
  for (sh_iter64 i=0;i<max_lapic_id+1;i++) {
    sh_lapic_DEVICE *lapic_dev=lapic_dev_array[i];
    if (lapic_dev==SH_NULLPTR) continue;
    if (lapic_device_default==SH_NULLPTR) lapic_device_default=lapic_dev;
    lapic_found++;
  }
  if (lapic_found!=expected_lapic_found) return SH_STATUS_RESCAN_NEEDED;
  ap_bootstrap->lapic_base_pa=(sh_uint64)(sh_uint8*)lapic_device_default->base;
  sh_uint8 *cpu_struct_base_ptr=sh_malloc(sizeof(sh_ap_CPU_STRUCT)*(max_lapic_id+1));
  if (!cpu_struct_base_ptr) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  sh_ap_CPU_STRUCT *cpu_struct_base=(sh_ap_CPU_STRUCT*)cpu_struct_base_ptr;
  ap_bootstrap->cpu_struct_base_area=(sh_uint64)cpu_struct_base_ptr;
  sh_uint8 *tss_array_ptr=sh_malloc(sizeof(sh_tss_TSS)*(lapic_found-1));
  sh_tss_TSS *tss_array=(sh_tss_TSS*)tss_array_ptr;
  sh_uint64 cpu_count=0;
  sh_uint64 cpu_index=1;
  for (sh_iter64 i=0;i<max_lapic_id+1;i++) {
    sh_lapic_DEVICE *lapic_dev=lapic_dev_array[i];
    if (lapic_dev==SH_NULLPTR) continue;
    if (lapic_dev->apic_id==current_lapic_id) continue;
    sh_uint8 *c_entry_point_stack=sh_malloc(64*SH_PAGE_SIZE);
    if (!c_entry_point_stack) return SH_STATUS_ERROR_NULLPTR_RETURNED;
    cpu_struct_base[i].c_entry_point_stack_top_va=(sh_uint64)c_entry_point_stack+64*SH_PAGE_SIZE;
    cpu_struct_base[i].lapic_id=lapic_dev->apic_id;
    status=sh_tss_fill_tss(&tss_array[cpu_count]);
    if (sh_status_error(status)) return status;
    cpu_struct_base[i].tss_selector=(sh_uint16)(0x28+cpu_count*sizeof(sh_gdt_GDT_ENTRY_128));
    cpu_struct_base[i].per_cpu=(sh_ap_PER_CPU*)sh_malloc(sizeof(sh_ap_PER_CPU));
    if (!cpu_struct_base[i].per_cpu) return SH_STATUS_ERROR_NULLPTR_RETURNED;
    cpu_struct_base[i].per_cpu->cpu_id=cpu_index;
    cpu_struct_base[i].per_cpu->bytes_outputed=0;
    cpu_struct_base[i].per_cpu->timer_state=SH_FALSE;
    cpu_struct_base[i].per_cpu->temp_buffer=sh_malloc(256);
    if (!cpu_struct_base[i].per_cpu->temp_buffer) return SH_STATUS_ERROR_NULLPTR_RETURNED;
    cpu_struct_base[i].per_cpu->temp_buffer_size=256;
    cpu_count++;
    cpu_index++;
  }
  sh_uint8 *global_gdt_ptr=sh_malloc(sizeof(sh_gdt_GDT_AP));
  sh_gdt_GDT_AP *global_gdt=(sh_gdt_GDT_AP*)global_gdt_ptr;
  status=sh_gdt_fill_gdt_ap(tss_array,cpu_count,global_gdt);
  if (sh_status_error(status)) return status;
  ap_bootstrap->shared_gdt_64_va=(sh_uint64)global_gdt;
  ap_bootstrap->gdt_64_gdtr=sh_gdt_make_gdtr_ap(global_gdt);
  ap_bootstrap->idt=idt;
  sh_smp_set_cpu_count((sh_int16)cpu_index);
  is_ap_started=sh_malloc(cpu_index);
  if (!is_ap_started) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  sh_mem_set_8(is_ap_started,SH_AP_STATE_PENDING,cpu_index);
  return SH_STATUS_SUCCESS;
}
void sh_ap_entry_point() {
  sh_lapic_DEVICE *lapic_dev=sh_lapic_get_dev_apic_id(sh_smp_gs_base()->lapic_id);
  if (!lapic_dev) {
    sh_print(SH_FATAL,"Couldn't obtain LAPIC of this AP.");
    is_ap_started[sh_smp_gs_base()->per_cpu->cpu_id]=SH_AP_STATE_NO_LAPIC;
    while (SH_TRUE) {};
  }
  SH_STATUS status=sh_lapic_init_dev(0xFF,lapic_dev);
  if (sh_status_error(status)) {
    sh_print(SH_FATAL,"Couldn't initialize LAPIC.");
    is_ap_started[sh_smp_gs_base()->per_cpu->cpu_id]=SH_AP_STATE_NO_LAPIC;
    while (SH_TRUE) {};
  }
  sh_idt_load_idtr(ap_bootstrap_struct->idt);
  sh_asm_sti();
  sh_printf(SH_LOG,"Thread with lapic_id=%4u and cpu_id=%8u successfully reached C entry point.\n",sh_smp_gs_base()->lapic_id,sh_smp_gs_base()->per_cpu->cpu_id);
  sh_mb();
  is_ap_started[sh_smp_gs_base()->per_cpu->cpu_id]=SH_AP_STATE_OK;
  sh_mb();
  while (SH_TRUE) {};
}
SH_STATUS sh_ap_start_ap_boot_procedure(sh_conf_BOOT_CONFIG *boot_config) {
  SH_STATUS status=sh_memory_identity_map(SH_AP_AP_TRAMPOLINE,1,SH_PAGE_RW | SH_PAGE_PRESENT);
  if (sh_status_error(status)) return status;
  sh_mem_copy((sh_uint8*)SH_AP_AP_TRAMPOLINE,(sh_uint8*)SH_AP_TRAMPOLINE_PAYLOAD,4096);
  sh_lapic_DEVICE *lapic_dev=sh_lapic_get_dev_apic_id((sh_uint64)sh_lapic_get_current_core());
  if (lapic_dev==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  bsp_cpu_struct.c_entry_point_stack_top_va=0;
  bsp_cpu_struct.lapic_id=(sh_uint32)sh_lapic_get_current_core();
  bsp_cpu_struct.tss_selector=0;
  bsp_per_cpu.cpu_id=0;
  bsp_per_cpu.bytes_outputed=0;
  bsp_per_cpu.timer_state=SH_FALSE;
  bsp_per_cpu.temp_buffer=sh_malloc(256);
  if (!bsp_per_cpu.temp_buffer) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  bsp_per_cpu.temp_buffer_size=256;
  bsp_cpu_struct.per_cpu=(&bsp_per_cpu);
  sh_smp_write_gs_base((sh_uint64)&bsp_cpu_struct);
  is_ap_started[sh_smp_gs_base()->per_cpu->cpu_id]=SH_AP_STATE_OK;
  sh_log_debug("Trying sending INIT IPI...",SH_LOG_SOURCE_SMP);
  status=sh_lapic_send_ipi(lapic_dev,SH_LAPIC_IPI_TYPE_INIT,SH_LAPIC_IPI_DESTINATION_ALL_EXCLUDING_SELF,SH_LAPIC_IPI_NO_DESTINATION,0);
  if (sh_status_error(status)) return status;
  sh_log_debug("Sent INIT IPI. Waiting 10 ms.",SH_LOG_SOURCE_SMP);
  status=sh_lapic_timer_one_shot_us(lapic_dev,10000);
  if (sh_status_error(status)) return status;
  sh_log_debug("Trying sending first SIPI...",SH_LOG_SOURCE_SMP);
  status=sh_print_setup_ring_buffers_per_ap(boot_config);
  if (sh_status_error(status)) return status;
  status=sh_lapic_send_ipi(lapic_dev,SH_LAPIC_IPI_TYPE_STARTUP,SH_LAPIC_IPI_DESTINATION_ALL_EXCLUDING_SELF,SH_LAPIC_IPI_NO_DESTINATION,SH_AP_AP_TRAMPOLINE);
  if (sh_status_error(status)) return status;
  sh_log_disable_global_logging_ring();
  sh_print(SH_DEBUG,"Sent first SIPI.\n");
  while (SH_TRUE) {
    sh_mb();
    sh_int16 cpu_good=0;
    for (sh_int16 i=0;i<sh_smp_get_cpu_count();i++) {
      if (is_ap_started[i]==SH_AP_STATE_PENDING) continue;
      if (is_ap_started[i]==SH_AP_STATE_NO_LAPIC) {
        sh_printf(SH_CRITICAL,"AP %8u failed to set up his LAPIC.\n",i);
        return SH_STATUS_AP_NO_LAPIC;
      } else if (is_ap_started[i]==SH_AP_STATE_OK) {
        cpu_good++;
      }
    }
    if (cpu_good==sh_smp_get_cpu_count()) {
      return SH_STATUS_SUCCESS;
    }
  }
}