#include "hv/allocator.h"
#include "hv/arch.h"
#include "hv/config.h"
#include "hv/log.h"
#include "hv/mmio_policy.h"
#include "hv/panic.h"
#include "hv/stage2.h"
#include "hv/string.h"
#include "hv/sysreg.h"
#include "hv/trap.h"
#include "hv/uart.h"
#include "hv/vcpu.h"

#define S2_ENTRIES 512u
#define S2_L1_SHIFT 30u
#define S2_L2_SHIFT 21u
#define S2_L1_SIZE (1ull << S2_L1_SHIFT)
#define S2_L2_SIZE (1ull << S2_L2_SHIFT)
#define S2_DESC_VALID HV_BIT(0)
#define S2_DESC_TABLE HV_BIT(1)
#define S2_DESC_AF HV_BIT(10)
#define S2_DESC_SH_INNER (3ull << 8u)
#define S2_DESC_MEMATTR_NORMAL (0xfull << 2u)
#define S2_DESC_MEMATTR_DEVICE (0x0ull << 2u)
#define S2_DESC_S2AP_R (1ull << 6u)
#define S2_DESC_S2AP_W (1ull << 7u)
#define S2_DESC_XN (3ull << 53u)

static uint64_t g_l1[S2_ENTRIES] __attribute__((aligned(65536)));
static uint64_t g_l2_tables[4][S2_ENTRIES] __attribute__((aligned(4096)));
static uint32_t g_l2_used;

static uint64_t desc_addr(uint64_t addr)
{
    return addr & 0x0000fffffffff000ull;
}

void stage2_init(void)
{
    hv_memset(g_l1, 0, sizeof(g_l1));
    hv_memset(g_l2_tables, 0, sizeof(g_l2_tables));
    g_l2_used = 0;
}

static uint64_t *ensure_l2(uint64_t l1_index)
{
    if (l1_index >= S2_ENTRIES) {
        return (uint64_t *)0;
    }
    if ((g_l1[l1_index] & S2_DESC_VALID) != 0u) {
        return (uint64_t *)(uintptr_t)desc_addr(g_l1[l1_index]);
    }
    if (g_l2_used >= HV_ARRAY_SIZE(g_l2_tables)) {
        return (uint64_t *)0;
    }
    uint64_t *table = g_l2_tables[g_l2_used++];
    g_l1[l1_index] = ((uint64_t)(uintptr_t)table & 0x0000fffffffff000ull) |
                     S2_DESC_TABLE | S2_DESC_VALID;
    return table;
}

static uint64_t block_desc(paddr_t pa, enum stage2_mem_type type, uint32_t perms)
{
    uint64_t desc = (pa & 0x0000ffffffe00000ull) | S2_DESC_VALID | S2_DESC_AF;
    desc |= (type == S2_MEM_NORMAL) ? (S2_DESC_MEMATTR_NORMAL | S2_DESC_SH_INNER) :
                                      S2_DESC_MEMATTR_DEVICE;
    if ((perms & S2_PERM_R) != 0u) {
        desc |= S2_DESC_S2AP_R;
    }
    if ((perms & S2_PERM_W) != 0u) {
        desc |= S2_DESC_S2AP_W;
    }
    if ((perms & S2_PERM_X) == 0u) {
        desc |= S2_DESC_XN;
    }
    return desc;
}

static bool descriptor_valid(uint64_t desc)
{
    if ((desc & S2_DESC_VALID) == 0u) {
        return false;
    }
    if ((desc & S2_DESC_AF) == 0u) {
        return false;
    }
    if (!hv_is_aligned_u64(desc_addr(desc), S2_L2_SIZE)) {
        return false;
    }
    return true;
}

static uint64_t *lookup_l2(uint64_t l1_index)
{
    if (l1_index >= S2_ENTRIES || (g_l1[l1_index] & S2_DESC_VALID) == 0u) {
        return (uint64_t *)0;
    }
    if ((g_l1[l1_index] & S2_DESC_TABLE) == 0u) {
        return (uint64_t *)0;
    }
    return (uint64_t *)(uintptr_t)desc_addr(g_l1[l1_index]);
}

bool stage2_map_range(ipa_t ipa, paddr_t pa, uint64_t size, enum stage2_mem_type type, uint32_t perms)
{
    uint64_t ipa_end;
    uint64_t pa_end;

    if (size == 0u || perms == 0u) {
        return false;
    }
    if (!hv_is_aligned_u64(ipa, S2_L2_SIZE) || !hv_is_aligned_u64(pa, S2_L2_SIZE) ||
        !hv_is_aligned_u64(size, S2_L2_SIZE)) {
        return false;
    }
    if (hv_add_overflow_u64(ipa, size, &ipa_end) ||
        hv_add_overflow_u64(pa, size, &pa_end)) {
        return false;
    }
    if (ipa_end <= ipa || pa_end <= pa) {
        return false;
    }
    if (hv_range_overlaps(pa, size, (uint64_t)(uintptr_t)__image_start,
                          (uint64_t)((uintptr_t)__image_end - (uintptr_t)__image_start))) {
        return false;
    }

    for (uint64_t off = 0; off < size; off += S2_L2_SIZE) {
        uint64_t cur_ipa = ipa + off;
        uint64_t l1_index = (cur_ipa >> S2_L1_SHIFT) & 0x1ffu;
        uint64_t l2_index = (cur_ipa >> S2_L2_SHIFT) & 0x1ffu;
        uint64_t *l2 = ensure_l2(l1_index);
        if (l2 == (uint64_t *)0 || (l2[l2_index] & S2_DESC_VALID) != 0u) {
            return false;
        }
        uint64_t desc = block_desc(pa + off, type, perms);
        if (!descriptor_valid(desc)) {
            return false;
        }
        l2[l2_index] = desc;
    }
    arch_tlbi_vmalle1is();
    return true;
}

bool stage2_unmap_range(ipa_t ipa, uint64_t size)
{
    if (size == 0u || !hv_is_aligned_u64(ipa, S2_L2_SIZE) ||
        !hv_is_aligned_u64(size, S2_L2_SIZE)) {
        return false;
    }

    for (uint64_t off = 0; off < size; off += S2_L2_SIZE) {
        uint64_t cur_ipa = ipa + off;
        uint64_t l1_index = (cur_ipa >> S2_L1_SHIFT) & 0x1ffu;
        uint64_t l2_index = (cur_ipa >> S2_L2_SHIFT) & 0x1ffu;
        uint64_t *l2 = lookup_l2(l1_index);
        if (l2 == (uint64_t *)0 || (l2[l2_index] & S2_DESC_VALID) == 0u) {
            return false;
        }
    }

    for (uint64_t off = 0; off < size; off += S2_L2_SIZE) {
        uint64_t cur_ipa = ipa + off;
        uint64_t l1_index = (cur_ipa >> S2_L1_SHIFT) & 0x1ffu;
        uint64_t l2_index = (cur_ipa >> S2_L2_SHIFT) & 0x1ffu;
        uint64_t *l2 = lookup_l2(l1_index);
        l2[l2_index] = 0u;
    }
    arch_tlbi_vmalle1is();
    return true;
}

bool stage2_translate(ipa_t ipa, struct stage2_translation *out)
{
    uint64_t l1_index = (ipa >> S2_L1_SHIFT) & 0x1ffu;
    uint64_t l2_index = (ipa >> S2_L2_SHIFT) & 0x1ffu;
    uint64_t *l2 = lookup_l2(l1_index);
    uint64_t desc;

    if (out != (struct stage2_translation *)0) {
        hv_memset(out, 0, sizeof(*out));
    }
    if (l2 == (uint64_t *)0 || (l2[l2_index] & S2_DESC_VALID) == 0u) {
        return false;
    }

    desc = l2[l2_index];
    if (out != (struct stage2_translation *)0) {
        out->mapped = true;
        out->ipa_base = HV_ALIGN_DOWN(ipa, S2_L2_SIZE);
        out->pa_base = desc_addr(desc);
        out->size = S2_L2_SIZE;
        out->desc = desc;
        if ((desc & S2_DESC_S2AP_R) != 0u) {
            out->perms |= S2_PERM_R;
        }
        if ((desc & S2_DESC_S2AP_W) != 0u) {
            out->perms |= S2_PERM_W;
        }
        if ((desc & S2_DESC_XN) == 0u) {
            out->perms |= S2_PERM_X;
        }
        out->type = (desc & S2_DESC_MEMATTR_NORMAL) == S2_DESC_MEMATTR_NORMAL ?
                    S2_MEM_NORMAL : S2_MEM_DEVICE;
    }
    return true;
}

ipa_t stage2_ipa_from_abort(uint64_t far_el2, uint64_t hpfar_el2)
{
    return ((hpfar_el2 & 0x0000000ffffffff0ull) << 8u) | (far_el2 & 0xfffull);
}

static bool dfsc_is_translation(uint32_t dfsc)
{
    return dfsc >= 0x04u && dfsc <= 0x07u;
}

bool stage2_handle_abort(struct vcpu *vcpu, const struct trap_frame *frame, bool instruction)
{
    ipa_t ipa;
    ipa_t block_ipa;
    paddr_t block_pa;
    uint32_t dfsc;
    struct log_event ev;

    if (vcpu == (struct vcpu *)0 || frame == (const struct trap_frame *)0) {
        return false;
    }

    ipa = stage2_ipa_from_abort(frame->far_el2, frame->hpfar_el2);
    dfsc = (uint32_t)(frame->esr_el2 & 0x3fu);
    hv_memset(&ev, 0, sizeof(ev));
    ev.kind = LOG_ABORT;
    ev.vcpu = vcpu->id;
    ev.ec = esr_ec(frame->esr_el2);
    ev.iss = frame->esr_el2 & 0x01ffffffu;
    ev.elr = frame->elr_el2;
    ev.far = ipa;
    ev.hpfar = frame->hpfar_el2;
    ev.pstate = frame->spsr_el2;
    ev.reason = instruction ? 406u : 407u;

    const struct mmio_policy_region *mmio = mmio_policy_lookup(ipa);
    if (mmio != (const struct mmio_policy_region *)0) {
        ev.reason = 410u;
        ev.name = mmio->name;
        ev.action = (uint32_t)mmio->action;
        if (mmio_policy_emulate(vcpu, ipa, frame->esr_el2)) {
            ev.reason = 411u;
            log_event_commit(&ev);
            arch_advance_guest_pc(vcpu);
            return true;
        }
        log_event_commit(&ev);
        return false;
    }

    if (!dfsc_is_translation(dfsc) ||
        ipa < CONFIG_GUEST_IPA_BASE ||
        ipa >= CONFIG_GUEST_IPA_BASE + CONFIG_GUEST_RAM_SIZE) {
        log_event_commit(&ev);
        return false;
    }

    block_ipa = HV_ALIGN_DOWN(ipa, S2_L2_SIZE);
    block_pa = CONFIG_GUEST_RAM_BASE + (block_ipa - CONFIG_GUEST_IPA_BASE);
    if (!stage2_map_range(block_ipa, block_pa, S2_L2_SIZE, S2_MEM_NORMAL,
                          S2_PERM_R | S2_PERM_W | S2_PERM_X)) {
        log_event_commit(&ev);
        return false;
    }

    ev.reason = instruction ? 408u : 409u;
    ev.result = block_pa;
    log_event_commit(&ev);
    return true;
}

uint64_t stage2_vttbr(void)
{
    return (uint64_t)(uintptr_t)g_l1;
}

uint64_t stage2_vtcr(void)
{
    uint64_t vtcr = 0;
    vtcr |= 24ull;
    vtcr |= (1ull << 6u);
    vtcr |= (1ull << 8u);
    vtcr |= (1ull << 10u);
    vtcr |= (3ull << 12u);
    vtcr |= (2ull << 16u);
    return vtcr;
}

void stage2_dump(void)
{
    uart_puts("stage2 root=");
    uart_put_hex64(stage2_vttbr());
    uart_puts(" vtcr=");
    uart_put_hex64(stage2_vtcr());
    uart_puts(" l2_tables=");
    uart_put_dec64(g_l2_used);
    uart_puts("\n");
    for (uint32_t l1 = 0; l1 < S2_ENTRIES; l1++) {
        if ((g_l1[l1] & S2_DESC_VALID) == 0u) {
            continue;
        }
        uint64_t *l2 = (uint64_t *)(uintptr_t)desc_addr(g_l1[l1]);
        for (uint32_t l2i = 0; l2i < S2_ENTRIES; l2i++) {
            if ((l2[l2i] & S2_DESC_VALID) != 0u) {
                uart_puts("mapping ipa=");
                uart_put_hex64(((uint64_t)l1 << S2_L1_SHIFT) | ((uint64_t)l2i << S2_L2_SHIFT));
                uart_puts(" desc=");
                uart_put_hex64(l2[l2i]);
                uart_puts("\n");
            }
        }
    }
}

bool stage2_selftest(void)
{
    if (!hv_is_aligned_u64((uint64_t)(uintptr_t)g_l1, 65536u)) {
        return false;
    }
    for (uint32_t i = 0; i < S2_ENTRIES; i++) {
        if ((g_l1[i] & S2_DESC_VALID) != 0u && (g_l1[i] & S2_DESC_TABLE) == 0u) {
            return false;
        }
        if ((g_l1[i] & S2_DESC_VALID) != 0u) {
            uint64_t *l2 = (uint64_t *)(uintptr_t)desc_addr(g_l1[i]);
            for (uint32_t j = 0; j < S2_ENTRIES; j++) {
                if ((l2[j] & S2_DESC_VALID) != 0u && !descriptor_valid(l2[j])) {
                    return false;
                }
            }
        }
    }
    return stage2_vttbr() == (uint64_t)(uintptr_t)g_l1;
}