bits 16

section .rodata

global smp_trampoline_start
smp_trampoline_start:
    cli
    cld

    mov ebx, cs
    shl ebx, 4

    o32 lidt [cs:(invalid_idt - smp_trampoline_start)]
    o32 lgdt [cs:(passed_info.gdtr - smp_trampoline_start)]

    lea eax, [ebx + (.mode32 - smp_trampoline_start)]
    mov [cs:(.farjmp_off - smp_trampoline_start)], eax

    mov eax, 0x00000011
    mov cr0, eax
    o32 jmp far [cs:(.farjmp - smp_trampoline_start)]

  .farjmp:
    .farjmp_off: dd 0
    .farjmp_seg: dd 0x18

    bits 32
  .mode32:
    mov ax, 0x20
    mov ds, ax
    mov es, ax
    mov fs, ax
    mov gs, ax
    mov ss, ax

    xor eax, eax
    lldt ax

    xor eax, eax
    mov cr4, eax

    mov esi, ebx
    mov eax, 1
    xor ecx, ecx
    cpuid
    test edx, 1 << 16
    jz .no_pat
    mov ecx, 0x277
    mov eax, 0x00070406
    mov edx, 0x00000105
    wrmsr
  .no_pat:
    mov ebx, esi

    mov ecx, 0x1b
    rdmsr
    test eax, (1 << 10)
    jz .write_apic_msr

    ; Check if target also has x2APIC
    test dword [ebx + (passed_info.bsp_apic_addr_msr_lo - smp_trampoline_start)], (1 << 10)
    jnz .write_apic_msr

    ; AP is x2APIC but target is xAPIC: go through disabled state
    btr eax, 11
    btr eax, 10
    wrmsr

  .write_apic_msr:
    mov eax, [ebx + (passed_info.bsp_apic_addr_msr_lo - smp_trampoline_start)]
    mov edx, [ebx + (passed_info.bsp_apic_addr_msr_hi - smp_trampoline_start)]
    bts eax, 11
    btr eax, 8
    wrmsr
    mov esp, [ebx + (passed_info.temp_stack - smp_trampoline_start)]

    mov eax, cr4
    bts eax, 5
    mov cr4, eax

    ; Set EFER (LME + NX if available), all other bits cleared
    mov ecx, 0xc0000080
    xor edx, edx
    mov eax, 1 << 8
    test dword [ebx + (passed_info.target_mode - smp_trampoline_start)], (1 << 3)
    jz .no_nx
    or eax, 1 << 11
  .no_nx:
    wrmsr

    test dword [ebx + (passed_info.target_mode - smp_trampoline_start)], (1 << 1)
    jz .no5lv

    mov eax, cr4
    bts eax, 12
    mov cr4, eax

  .no5lv:
    mov eax, dword [ebx + (passed_info.pagemap - smp_trampoline_start)]
    mov cr3, eax

    ; Enable CR0.PG (and WP if requested)
    mov eax, cr0
    test dword [ebx + (passed_info.target_mode - smp_trampoline_start)], (1 << 4)
    jz .no_wp
    bts eax, 16
  .no_wp:
    bts eax, 31
    mov cr0, eax

%ifdef IA32_TARGET
    ; Synchronise MTRRs with BSP
    call [ebx + (passed_info.mtrr_restore - smp_trampoline_start)]

    ; Configure local APIC handoff state (if pointer is set)
    mov eax, dword [ebx + (passed_info.lapic_setup - smp_trampoline_start)]
    test eax, eax
    jz .skip_lapic_setup32
    call eax
  .skip_lapic_setup32:
%endif

    lea eax, [ebx + (.mode64 - smp_trampoline_start)]
    push 0x28
    push eax
    retf

    bits 64
  .mode64:
    mov ax, 0x30
    mov ds, ax
    mov es, ax
    mov fs, ax
    mov gs, ax
    mov ss, ax

    mov ebx, ebx
    mov rax, qword [rbx + (passed_info.hhdm - smp_trampoline_start)]
    add qword [rbx + (passed_info.gdtr - smp_trampoline_start) + 2], rax
    lgdt [rbx + (passed_info.gdtr - smp_trampoline_start)]

    ; Clear TSS busy bit and load TR
    mov rcx, [rbx + (passed_info.gdtr - smp_trampoline_start) + 2]
    mov byte [rcx + 0x3d], 0x89
    mov ecx, 0x38
    ltr cx

    lea rax, [rax + rbx + (parking64 - smp_trampoline_start)]

    jmp rax

bits 64
parking64:
    mov ebx, ebx

%ifdef X86_64_TARGET
    ; Synchronise MTRRs with BSP
    call [rbx + (passed_info.mtrr_restore - smp_trampoline_start)]

    ; Configure local APIC handoff state (if pointer is set)
    mov rax, [rbx + (passed_info.lapic_setup - smp_trampoline_start)]
    test rax, rax
    jz .skip_lapic_setup64
    call rax
  .skip_lapic_setup64:
%endif

    mov edi, dword [rbx + (passed_info.smp_info_struct - smp_trampoline_start)]
    add rdi, qword [rbx + (passed_info.hhdm - smp_trampoline_start)]

    mov eax, 1
    xchg dword [rbx + (passed_info.booted_flag - smp_trampoline_start)], eax

    ; Check for MONITOR/MWAIT support
    mov eax, 1
    xor ecx, ecx
    cpuid
    test ecx, (1 << 3)
    jnz .monitor_spin

  .loop:
    mov rax, qword [rdi + 16]
    test rax, rax
    jnz .out
    pause
    jmp .loop

  .monitor_spin:
    mov rax, qword [rdi + 16]
    test rax, rax
    jnz .out
    lea rax, [rdi + 16]
    xor ecx, ecx
    xor edx, edx
    monitor
    mov rax, qword [rdi + 16]
    test rax, rax
    jnz .out
    xor eax, eax
    xor ecx, ecx
    mwait
    jmp .monitor_spin

  .out:
    ; Clear TLB
    mov rbx, cr3
    mov cr3, rbx

    ; Switch to new stack (HHDM address, safe after lower half unmap)
    mov rsp, qword [rdi + 8]

    ; Push fake return address
    push 0
    mov rsi, rsp

    ; Prepare iretq frame
    push 0x30
    push rsi
    push 0x2
    push 0x28
    push rax

    ; Zero out all GPRs (except rdi = mp_info pointer)
    xor eax, eax
    xor ebx, ebx
    xor ecx, ecx
    xor edx, edx
    xor esi, esi
    xor ebp, ebp
    xor r8d,  r8d
    xor r9d,  r9d
    xor r10d, r10d
    xor r11d, r11d
    xor r12d, r12d
    xor r13d, r13d
    xor r14d, r14d
    xor r15d, r15d

    iretq

invalid_idt:
    times 2 dq 0

align 16
passed_info:
    .booted_flag db 0
    .target_mode db 0
    .pagemap dd 0
    .smp_info_struct dd 0
    .gdtr:
        dw 0
        dq 0
    .hhdm:
        dq 0
    .bsp_apic_addr_msr_lo:
        dd 0
    .bsp_apic_addr_msr_hi:
        dd 0
    .mtrr_restore:
        dq 0
    .temp_stack:
        dq 0
    .lapic_setup:
        dq 0

smp_trampoline_end:

global smp_trampoline_size
smp_trampoline_size dq smp_trampoline_end - smp_trampoline_start

section .note.GNU-stack noalloc noexec nowrite progbits