#include "hv/arch.h"
#include "hv/config.h"
#include "hv/el2_mmu.h"
#include "hv/string.h"
#include "hv/uart.h"

#define EL2_PT_ENTRIES 512u
#define EL2_L1_SHIFT 30u
#define EL2_L2_SHIFT 21u
#define EL2_L2_SIZE (1ull << EL2_L2_SHIFT)
#define EL2_DESC_VALID HV_BIT(0)
#define EL2_DESC_TABLE HV_BIT(1)
#define EL2_DESC_AF HV_BIT(10)
#define EL2_DESC_SH_INNER (3ull << 8u)
#define EL2_DESC_ATTR_NORMAL (1ull << 2u)
#define EL2_DESC_ATTR_DEVICE (0ull << 2u)
#define EL2_DESC_PXN HV_BIT(53)
#define EL2_DESC_UXN HV_BIT(54)
#define EL2_MAIR ((0xffull << 8u) | 0x04ull)

static uint64_t g_el2_l1[EL2_PT_ENTRIES] __attribute__((aligned(4096)));
static uint64_t g_el2_l2[2][EL2_PT_ENTRIES] __attribute__((aligned(4096)));
static uint64_t g_guard_low;
static uint64_t g_guard_high;
static bool g_ready;
static bool g_enabled;

static uint64_t block_desc(uint64_t pa, bool device, bool executable)
{
    uint64_t desc = (pa & 0x0000ffffffe00000ull) | EL2_DESC_VALID | EL2_DESC_AF;
    desc |= device ? EL2_DESC_ATTR_DEVICE : (EL2_DESC_ATTR_NORMAL | EL2_DESC_SH_INNER);
    if (!executable) {
        desc |= EL2_DESC_PXN | EL2_DESC_UXN;
    }
    return desc;
}

static bool install_block(uint64_t pa, bool device, bool executable)
{
    uint64_t l1 = (pa >> EL2_L1_SHIFT) & 0x1ffu;
    uint64_t l2 = (pa >> EL2_L2_SHIFT) & 0x1ffu;
    if (l1 >= HV_ARRAY_SIZE(g_el2_l2)) {
        return false;
    }
    g_el2_l1[l1] = ((uint64_t)(uintptr_t)g_el2_l2[l1] & 0x0000fffffffff000ull) |
                   EL2_DESC_TABLE | EL2_DESC_VALID;
    g_el2_l2[l1][l2] = block_desc(pa, device, executable);
    return true;
}

void el2_mmu_init_plan(void)
{
    uint64_t image_start = HV_ALIGN_DOWN((uint64_t)(uintptr_t)__image_start, EL2_L2_SIZE);
    uint64_t image_end = HV_ALIGN_UP((uint64_t)(uintptr_t)__image_end, EL2_L2_SIZE);

    hv_memset(g_el2_l1, 0, sizeof(g_el2_l1));
    hv_memset(g_el2_l2, 0, sizeof(g_el2_l2));
    g_guard_low = image_start >= EL2_L2_SIZE ? image_start - EL2_L2_SIZE : 0u;
    g_guard_high = image_end;
    g_ready = true;
    g_enabled = false;

    for (uint64_t pa = image_start; pa < image_end; pa += EL2_L2_SIZE) {
        if (!install_block(pa, false, true)) {
            g_ready = false;
        }
    }
    for (uint64_t pa = CONFIG_GUEST_RAM_BASE;
         pa < CONFIG_GUEST_RAM_BASE + CONFIG_GUEST_RAM_SIZE;
         pa += EL2_L2_SIZE) {
        if (!install_block(pa, false, false)) {
            g_ready = false;
        }
    }
    if (!install_block(CONFIG_UART_BASE, true, false) ||
        !install_block(CONFIG_GICD_BASE, true, false) ||
        !install_block(CONFIG_GICR_BASE, true, false) ||
        !install_block(0x0a000000ull, true, false)) {
        g_ready = false;
    }
}

uint64_t el2_mmu_root(void)
{
    return (uint64_t)(uintptr_t)g_el2_l1;
}

uint64_t el2_mmu_tcr(void)
{
    uint64_t tcr = 0u;
    tcr |= 25ull;
    tcr |= (1ull << 8u);
    tcr |= (1ull << 10u);
    tcr |= (3ull << 12u);
    tcr |= (2ull << 16u);
    return tcr;
}

uint64_t el2_mmu_mair(void)
{
    return EL2_MAIR;
}

void el2_mmu_enable(void)
{
    if (!g_ready) {
        return;
    }
    arch_enable_el2_mmu(el2_mmu_root(), el2_mmu_tcr(), el2_mmu_mair());
    g_enabled = true;
}

void el2_mmu_dump(void)
{
    uart_puts("el2mmu planned=");
    uart_puts(g_ready ? "yes" : "no");
    uart_puts(" enabled=");
    uart_puts(g_enabled ? "yes" : "no");
    uart_puts(" root=");
    uart_put_hex64(el2_mmu_root());
    uart_puts(" tcr=");
    uart_put_hex64(el2_mmu_tcr());
    uart_puts(" mair=");
    uart_put_hex64(el2_mmu_mair());
    uart_puts(" guard_low=");
    uart_put_hex64(g_guard_low);
    uart_puts(" guard_high=");
    uart_put_hex64(g_guard_high);
    uart_puts("\n");
}

bool el2_mmu_selftest(void)
{
    return g_ready &&
           g_enabled &&
           hv_is_aligned_u64(el2_mmu_root(), CONFIG_PAGE_SIZE) &&
           (g_guard_high == 0u || (g_guard_high & (EL2_L2_SIZE - 1u)) == 0u);
}