#include <efi.h>
#include <lib.h>
#include <efilib.h>
#define _EFI_MP_
#include "efi-mp.h"

VOID EFIAPI SMP_function(IN VOID* buf) {
    EFI_GUID mpspGuid = EFI_MP_SERVICES_PROTOCOL_GUID;
    EFI_MP_SERVICES_PROTOCOL *mp = 0;

    EFI_STATUS status = uefi_call_wrapper(BS->LocateProtocol, 3, &mpspGuid, NULL, (void**)&mp);
    if(EFI_ERROR(status) || !mp) {
        Print(L"[-] Failed to locate MP protocol in SMP\n");
        return;
    }

    int id = 0;
    status = uefi_call_wrapper(mp->WhoAmI, 2, mp, &id);
    if(EFI_ERROR(status)) {
        Print(L"[-] Failed to get ID in SMP\n");
        return;
    }
    for (int i = 0; i < 100; i++) {
        Print(L"Hello from SMP: %d\n", id);
    }
}

EFI_STATUS
EFIAPI
efi_main (EFI_HANDLE image, EFI_SYSTEM_TABLE *SystemTable)
{
    UINTN argc;
    CHAR16 **argv;
    EFI_STATUS status;

    InitializeLib(image, SystemTable);

    EFI_EVENT Event = 0;
    EFI_GUID mpspGuid = EFI_MP_SERVICES_PROTOCOL_GUID;
    EFI_MP_SERVICES_PROTOCOL *mp = 0;
    UINT8 pibuffer[100];
    EFI_PROCESSOR_INFORMATION *pibuf=(EFI_PROCESSOR_INFORMATION*)pibuffer;
    UINTN NumberOfProcessors, NumberOfEnabledProcessors;
    UINTN BSP_proc = -1;

    status = uefi_call_wrapper(BS->LocateProtocol, 3, &mpspGuid, NULL, (void**)&mp);
    if(EFI_ERROR(status) || !mp) {
        Print(L"[-] Failed to locate MP protocol\n");
        return EFI_SUCCESS;
    }
    status = uefi_call_wrapper(mp->GetNumberOfProcessors, 3, mp, &NumberOfProcessors, &NumberOfEnabledProcessors);
    if(EFI_ERROR(status) || !mp) {
        Print(L"[-] Failed to get number of cores\n");
        return EFI_SUCCESS;
    }
    Print(L"SMP numcores %d %d\n", NumberOfProcessors, NumberOfEnabledProcessors);
    if (NumberOfEnabledProcessors < 2) {
        Print(L"[-] Not enough cores\n");
        return EFI_SUCCESS;
    }

    // start APs
    status = uefi_call_wrapper(BS->CreateEvent, 5, 0, TPL_NOTIFY, NULL, NULL, &Event);
    if(EFI_ERROR(status) || !Event) {
        Print(L"[-] Failed to create event\n");
        return EFI_SUCCESS;
    }
    for(int i=0; i<NumberOfEnabledProcessors; i++) {
        status = uefi_call_wrapper(mp->GetProcessorInfo, 5, mp, i, pibuf);
        if(EFI_ERROR(status)) {
            Print(L"[-] Failed to get processor info\n");
            return EFI_SUCCESS;
        }
        if(pibuf->StatusFlag & PROCESSOR_AS_BSP_BIT) {
            BSP_proc = i;
        }
        if((pibuf->StatusFlag & PROCESSOR_ENABLED_BIT) == 0) {
            Print(L"Processor %d not enabled\n", i);
            return EFI_SUCCESS;
        }
        if((pibuf->StatusFlag & PROCESSOR_HEALTH_STATUS_BIT) == 0) {
            Print(L"Processor %d not healthy\n", i);
            return EFI_SUCCESS;
        }
    }
    if(BSP_proc == -1) {
        Print(L"[-] Failed to get BSP\n");
        return EFI_SUCCESS;
    }
    Print(L"spawning %d %lx\n", BSP_proc+1, Event);

    // status = uefi_call_wrapper(mp->StartupAllAPs, 7, mp, SMP_function, FALSE, Event, 0, (VOID*)NULL, NULL);
    status = uefi_call_wrapper(mp->StartupThisAP, 7, mp, SMP_function, BSP_proc+1, Event, 0, (VOID*) NULL, NULL);

    if(EFI_ERROR(status)) {
        Print(L"[-] Failed to spawn SMP\n");
        CHAR16 ErrorString[0x100] = {0};
        StatusToString(ErrorString, status);
        Print(L"ERROR: %s\n", ErrorString);
        return EFI_SUCCESS;
    }

    int id = 0;
    status = uefi_call_wrapper(mp->WhoAmI, 2, mp, &id);
    if(EFI_ERROR(status)) {
        Print(L"[-] Failed to get ID in BSP\n");
        return EFI_SUCCESS;
    }
    for (int i = 0; i < 100; i++) {
        Print(L"Hello from BSP: %d\n", id);
    }

  return EFI_SUCCESS;
}