#include "hv/arch.h"
#include "hv/config.h"
#include "hv/mmio.h"
#include "hv/panic.h"
#include "hv/uart.h"
#include "hv/vcpu.h"

uint64_t arch_current_el(void)
{
    uint64_t v;
    __asm__ volatile("mrs %0, CurrentEL" : "=r"(v));
    return (v >> 2u) & 0x3u;
}

uint64_t arch_mpidr_el1(void)
{
    uint64_t v;
    __asm__ volatile("mrs %0, mpidr_el1" : "=r"(v));
    return v;
}

uint64_t arch_cntpct_el0(void)
{
    uint64_t v;
    __asm__ volatile("mrs %0, cntpct_el0" : "=r"(v));
    return v;
}

uint64_t arch_cntfrq_el0(void)
{
    uint64_t v;
    __asm__ volatile("mrs %0, cntfrq_el0" : "=r"(v));
    return v;
}

uint64_t arch_vectors_base(void)
{
    return (uint64_t)(uintptr_t)__vectors_el2;
}

bool arch_vectors_aligned(void)
{
    return hv_is_aligned_u64(arch_vectors_base(), 2048u);
}

uint64_t arch_irq_save(void)
{
    uint64_t flags;
    __asm__ volatile("mrs %0, daif\nmsr daifset, #0xf\nisb" : "=r"(flags) :: "memory");
    return flags;
}

void arch_irq_restore(uint64_t flags)
{
    __asm__ volatile("msr daif, %0\nisb" :: "r"(flags) : "memory");
}

void arch_context_barrier(void)
{
    __asm__ volatile("dsb ish\nisb" ::: "memory");
}

void arch_dump_el2_state(void)
{
    uint64_t current_el;
    uint64_t hcr;
    uint64_t vtcr;
    uint64_t vttbr;
    uint64_t vbar;
    uint64_t sctlr;
    uint64_t mair;
    uint64_t sp;

    __asm__ volatile(
        "mrs %0, CurrentEL\n"
        "mrs %1, hcr_el2\n"
        "mrs %2, vtcr_el2\n"
        "mrs %3, vttbr_el2\n"
        "mrs %4, vbar_el2\n"
        "mrs %5, sctlr_el2\n"
        "mrs %6, mair_el2\n"
        "mov %7, sp\n"
        : "=r"(current_el), "=r"(hcr), "=r"(vtcr), "=r"(vttbr),
          "=r"(vbar), "=r"(sctlr), "=r"(mair), "=r"(sp));

    uart_puts("cpu current_el=");
    uart_put_hex64(current_el);
    uart_puts(" hcr_el2=");
    uart_put_hex64(hcr);
    uart_puts(" vtcr_el2=");
    uart_put_hex64(vtcr);
    uart_puts(" vttbr_el2=");
    uart_put_hex64(vttbr);
    uart_puts("\n");
    uart_puts("cpu vbar_el2=");
    uart_put_hex64(vbar);
    uart_puts(" sctlr_el2=");
    uart_put_hex64(sctlr);
    uart_puts(" mair_el2=");
    uart_put_hex64(mair);
    uart_puts(" sp=");
    uart_put_hex64(sp);
    uart_puts("\n");
}

void arch_install_vectors(void)
{
    if (!arch_vectors_aligned()) {
        panic("VBAR_EL2 vector table is not 2KB aligned");
    }
    __asm__ volatile("msr vbar_el2, %0\nisb" :: "r"(arch_vectors_base()) : "memory");
}

void arch_tlbi_vmalle1is(void)
{
    __asm__ volatile("dsb ish\n tlbi vmalle1is\n dsb ish\n isb" ::: "memory");
}

void arch_enable_el2_mmu(uint64_t ttbr0, uint64_t tcr, uint64_t mair)
{
    uint64_t sctlr;

    __asm__ volatile(
        "dsb sy\n"
        "msr ttbr0_el2, %0\n"
        "msr tcr_el2, %1\n"
        "msr mair_el2, %2\n"
        "isb\n"
        "tlbi alle2\n"
        "dsb sy\n"
        "isb\n"
        :: "r"(ttbr0), "r"(tcr), "r"(mair)
        : "memory");

    __asm__ volatile("mrs %0, sctlr_el2" : "=r"(sctlr));
    sctlr |= HV_BIT(0) | HV_BIT(2) | HV_BIT(12);
    __asm__ volatile(
        "msr sctlr_el2, %0\n"
        "isb\n"
        :: "r"(sctlr)
        : "memory");
}

void arch_el2_configure(uint64_t vttbr, uint64_t vtcr)
{
    const uint64_t mair = 0x000000000000ff00ull | 0x0000000000000004ull;
    uint64_t hcr = 0;

    hcr |= HV_BIT(31);
    hcr |= HV_BIT(0);
#if !CONFIG_BOOT_LINUX
    hcr |= HV_BIT(26);
#endif
    hcr |= HV_BIT(22);
    hcr |= HV_BIT(21);
    hcr |= HV_BIT(19);
#if !CONFIG_BOOT_LINUX
    hcr |= HV_BIT(20);
    hcr |= HV_BIT(18) | HV_BIT(17) | HV_BIT(16) | HV_BIT(15);
#endif
#if CONFIG_ENABLE_WFI_TRAPS
    hcr |= HV_BIT(14) | HV_BIT(13);
#endif

    __asm__ volatile(
        "msr mair_el2, %0\n"
        "msr vtcr_el2, %1\n"
        "msr vttbr_el2, %2\n"
        "msr mdcr_el2, %3\n"
        "msr cptr_el2, %4\n"
        "msr cnthctl_el2, %5\n"
        "msr hcr_el2, %6\n"
        "isb\n"
        :: "r"(mair), "r"(vtcr), "r"(vttbr),
           "r"(0ull), "r"(0x33ffull), "r"(3ull), "r"(hcr)
        : "memory");
}

void arch_prepare_guest_entry(struct vcpu *vcpu)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    __asm__ volatile(
        "msr sctlr_el1, %0\n"
        "msr ttbr0_el1, %1\n"
        "msr ttbr1_el1, %2\n"
        "msr tcr_el1, %3\n"
        "msr mair_el1, %4\n"
        "msr vbar_el1, %5\n"
        "msr cntkctl_el1, %6\n"
        "isb\n"
        :: "r"(vcpu->sysregs.sctlr_el1), "r"(vcpu->sysregs.ttbr0_el1),
           "r"(vcpu->sysregs.ttbr1_el1), "r"(vcpu->sysregs.tcr_el1),
           "r"(vcpu->sysregs.mair_el1), "r"(vcpu->sysregs.vbar_el1),
           "r"(vcpu->sysregs.cntkctl_el1)
        : "memory");
}

void arch_capture_el1_sysregs(struct vcpu *vcpu)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    __asm__ volatile(
        "mrs %0, vbar_el1\n"
        "mrs %1, elr_el1\n"
        "mrs %2, spsr_el1\n"
        "mrs %3, esr_el1\n"
        "mrs %4, far_el1\n"
        : "=r"(vcpu->sysregs.vbar_el1), "=r"(vcpu->sysregs.elr_el1),
          "=r"(vcpu->sysregs.spsr_el1), "=r"(vcpu->sysregs.esr_el1),
          "=r"(vcpu->sysregs.far_el1));
}

void arch_sync_el1_sysregs(const struct vcpu *vcpu)
{
    BUG_ON(vcpu == (const struct vcpu *)0);
    __asm__ volatile(
        "dsb sy\n"
        "msr mair_el1, %0\n"
        "msr tcr_el1, %1\n"
        "msr ttbr0_el1, %2\n"
        "msr ttbr1_el1, %3\n"
        "msr vbar_el1, %4\n"
        "msr cntkctl_el1, %5\n"
        "msr cpacr_el1, %6\n"
        "isb\n"
        "msr sctlr_el1, %7\n"
        "isb\n"
        :: "r"(vcpu->sysregs.mair_el1), "r"(vcpu->sysregs.tcr_el1),
           "r"(vcpu->sysregs.ttbr0_el1), "r"(vcpu->sysregs.ttbr1_el1),
           "r"(vcpu->sysregs.vbar_el1), "r"(vcpu->sysregs.cntkctl_el1),
           "r"(vcpu->sysregs.cpacr_el1), "r"(vcpu->sysregs.sctlr_el1)
        : "memory");
}

void arch_advance_guest_pc(struct vcpu *vcpu)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    vcpu->pc += 4u;
}

void arch_inject_guest_undef(struct vcpu *vcpu)
{
    const uint64_t current_el_spx_sync = 0x200u;
    const uint64_t pstate_el1h_daif = 0x3c5u;

    BUG_ON(vcpu == (struct vcpu *)0);
    arch_capture_el1_sysregs(vcpu);
    if (!hv_is_aligned_u64(vcpu->sysregs.vbar_el1, 2048u)) {
        vcpu->state = VCPU_PAUSED;
        return;
    }

    vcpu->sysregs.elr_el1 = vcpu->pc;
    vcpu->sysregs.spsr_el1 = vcpu->pstate;
    vcpu->sysregs.esr_el1 = 0u;
    vcpu->sysregs.far_el1 = 0u;
    vcpu->pc = vcpu->sysregs.vbar_el1 + current_el_spx_sync;
    vcpu->pstate = pstate_el1h_daif;
    vcpu->stats.injected_undefs++;

    __asm__ volatile(
        "msr elr_el1, %0\n"
        "msr spsr_el1, %1\n"
        "msr esr_el1, %2\n"
        "msr far_el1, %3\n"
        "isb\n"
        :: "r"(vcpu->sysregs.elr_el1), "r"(vcpu->sysregs.spsr_el1),
           "r"(vcpu->sysregs.esr_el1), "r"(vcpu->sysregs.far_el1)
        : "memory");
}

void arch_halt(void)
{
    for (;;) {
        __asm__ volatile("wfe");
    }
}