#include <lib/macros.aarch64_asm.h>

.set tpl_ap_el, -72
.set tpl_booted_flag, -64
.set tpl_hhdm_offset, -56
.set tpl_ttbr0, -48
.set tpl_ttbr1, -40
.set tpl_mair, -32
.set tpl_tcr, -24
.set tpl_sctlr, -16
.set tpl_info_struct, -8

.section .text

.global smp_trampoline_start
smp_trampoline_start:
    bl .L_entry
.L_entry:
    // Mask IRQs
    msr daifset, #0b1111

    // Address to next page (since our offsets into the boot data are negative)
    add x1, x30, #0xFFC

    ldr x0, [x1, tpl_info_struct]
    ldr x2, [x1, tpl_sctlr]
    ldr x3, [x1, tpl_mair]
    ldr x4, [x1, tpl_tcr]
    ldr x5, [x1, tpl_ttbr0]
    ldr x6, [x1, tpl_ttbr1]
    ldr x7, [x1, tpl_hhdm_offset]

    PICK_EL x8, 1f, 0f
0:
    // EL2 path - enable VHE and stay at EL2

    // Enable E2H if not already set
    mrs x8, hcr_el2
    orr x8, x8, #(1 << 34)
    msr hcr_el2, x8
    isb

    // Configure EL2 state for VHE
    msr spsel, #0

    mov x8, #3
    msr cnthctl_el2, x8
    msr cntvoff_el2, xzr

    // HCR: E2H + TGE + RW + SWIO
    mov x8, xzr
    orr x8, x8, #(1 << 34)
    orr x8, x8, #(1 << 31)
    orr x8, x8, #(1 << 27)
    orr x8, x8, #(1 << 1)
    msr hcr_el2, x8

    // Disable FP/SIMD/SVE (VHE CPTR_EL2 layout, CPACR-like)
    msr cptr_el2, xzr
    msr hstr_el2, xzr

    // Set up VBAR for page table switch (VHE redirects VBAR_EL1 to VBAR_EL2)
    adrp x8, 2f
    add x8, x8, :lo12:2f
    add x8, x8, x7
    msr vbar_el1, x8
    isb
    dsb sy
    isb

    // Switch page tables using VHE-redirected register names
    msr mair_el1, x3
    msr tcr_el1, x4
    msr ttbr0_el1, x5
    msr ttbr1_el1, x6
    msr sctlr_el1, x2
    isb
    dsb sy
    isb

    // Jump to the higher half mapping in case we didn't immediately crash
    br x8

1:
    // EL1 path
    msr spsel, #0

    // Switch to the new page tables

    // Point the EL1t handler to the continuation, such that after we page fault,
    // execution continues as expected.
    adrp x8, 2f
    add x8, x8, #:lo12:2f
    add x8, x8, x7
    msr vbar_el1, x8
    isb
    dsb sy
    isb

    // Switch the page table registers
    msr mair_el1, x3
    msr tcr_el1, x4
    msr ttbr0_el1, x5
    msr ttbr1_el1, x6
    msr sctlr_el1, x2
    isb
    dsb sy
    isb

    // Jump to the higher half mapping in case we didn't immediately crash
    br x8

// Alignment required by VBAR register
.align 12
2:
    // Zero out VBAR to avoid confusion
    msr vbar_el1, xzr

    // Disable FP/SIMD/SVE
    msr cpacr_el1, xzr

3:
    // Add HHDM offset to data pointer
    add x1, x1, x7

    // Report our EL to the BSP
    mrs x8, currentel
    lsr x8, x8, #2
    add x9, x1, tpl_ap_el
    stlr x8, [x9]

    // Notify BSP we are alive
    mov x8, #1
    add x9, x1, tpl_booted_flag
    stlr x8, [x9]

    // Wait for BSP to tell us where to go
    // Add HHDM offset to our info struct pointer
    add x0, x0, x7
    add x9, x0, #24
4:
    ldar x8, [x9]
    cbnz x8, 5f
    yield
    b 4b

5:
    msr elr_el1, x8

    msr spsel, #0
    ldr x8, [x0, #16]
    mov sp, x8

    // Enter kernel - determine SPSR from current EL
    mov x8, #0x3c4
    mrs x9, currentel
    cmp x9, #0b1000
    b.ne 9f
    mov x8, #0x3c8
9:
    msr spsr_el1, x8

    ZERO_REGS_EXCEPT_X0

    eret

smp_trampoline_end:

.section .rodata

.global smp_trampoline_size
smp_trampoline_size:
    .quad smp_trampoline_end - smp_trampoline_start

.section .note.GNU-stack,"",%progbits