#include "hv/arch.h"
#include "hv/config.h"
#include "hv/policy.h"
#include "hv/string.h"
#include "hv/timer.h"
#include "hv/uart.h"

#define KEY(op0, op1, crn, crm, op2) SYSREG_KEY_CONST((op0), (op1), (crn), (crm), (op2))
#define RO_ID(name, op0, op1, crn, crm, op2, val) \
    { KEY(op0, op1, crn, crm, op2), name, POLICY_CLASS_ID, POLICY_EMULATE_READ, POLICY_DENY_UNDEF, val, 0, 1, 100 }
#define SHADOW_RW(name, op0, op1, crn, crm, op2, field, mask) \
    { KEY(op0, op1, crn, crm, op2), name, POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, mask, 1, 101 }
#define SHADOW_THREAD(name, op0, op1, crn, crm, op2) \
    { KEY(op0, op1, crn, crm, op2), name, POLICY_CLASS_THREAD, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, UINT64_MAX, 1, 102 }
#define DENY_REG(name, op0, op1, crn, crm, op2, cls) \
    { KEY(op0, op1, crn, crm, op2), name, cls, POLICY_DENY_UNDEF, POLICY_DENY_UNDEF, 0, 0, 1, 200 }

static const struct policy_entry g_policy[] = {
    RO_ID("midr_el1", 3, 0, 0, 0, 0, 0x00000000410fd034ull),
    RO_ID("mpidr_el1", 3, 0, 0, 0, 5, 0x80000000ull),
    RO_ID("id_aa64pfr0_el1", 3, 0, 0, 4, 0, 0x0000000000000011ull),
    RO_ID("id_aa64pfr1_el1", 3, 0, 0, 4, 1, 0x0ull),
    RO_ID("id_aa64isar0_el1", 3, 0, 0, 6, 0, 0x0ull),
    RO_ID("id_aa64isar1_el1", 3, 0, 0, 6, 1, 0x0ull),
    RO_ID("id_aa64mmfr0_el1", 3, 0, 0, 7, 0, 0x0000000000001122ull),
    RO_ID("id_aa64mmfr1_el1", 3, 0, 0, 7, 1, 0x0ull),
    RO_ID("clidr_el1", 3, 1, 0, 0, 1, 0x0ull),
    { KEY(3, 2, 0, 0, 0), "csselr_el1", POLICY_CLASS_MISC, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, 0xfu, 1, 103 },
    RO_ID("ctr_el0", 3, 3, 0, 0, 1, 0x8444c004ull),
    RO_ID("dczid_el0", 3, 3, 0, 0, 7, 0x4ull),
    { KEY(3, 0, 1, 0, 0), "sctlr_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0x30d00800ull, 0x0000000030d0dfffull, 1, 104 },
    DENY_REG("actlr_el1", 3, 0, 1, 0, 1, POLICY_CLASS_MISC),
    { KEY(3, 0, 1, 0, 2), "cpacr_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, 0x00300000ull, 1, 105 },
    DENY_REG("hcr_el2", 3, 4, 1, 1, 0, POLICY_CLASS_EL2_EL3),
    DENY_REG("scr_el3", 3, 6, 1, 1, 0, POLICY_CLASS_EL2_EL3),
    { KEY(3, 0, 2, 0, 0), "ttbr0_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, 0x0000fffffffff000ull, 1, 106 },
    { KEY(3, 0, 2, 0, 1), "ttbr1_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, 0x0000fffffffff000ull, 1, 107 },
    { KEY(3, 0, 2, 0, 2), "tcr_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, 0x00000003ff3fffffull, 1, 108 },
    DENY_REG("vttbr_el2", 3, 4, 2, 1, 0, POLICY_CLASS_EL2_EL3),
    DENY_REG("pmcr_el0", 3, 3, 9, 12, 0, POLICY_CLASS_PMU),
    { KEY(3, 0, 10, 2, 0), "mair_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0xffull, UINT64_MAX, 1, 109 },
    { KEY(3, 0, 12, 0, 0), "vbar_el1", POLICY_CLASS_MMU, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, 0x0000fffffffff800ull, 1, 110 },
    DENY_REG("rvbar_el1", 3, 0, 12, 0, 1, POLICY_CLASS_MISC),
    SHADOW_THREAD("contextidr_el1", 3, 0, 13, 0, 1),
    SHADOW_THREAD("tpidr_el1", 3, 0, 13, 0, 4),
    SHADOW_THREAD("tpidrro_el0", 3, 3, 13, 0, 3),
    { KEY(3, 0, 14, 1, 0), "cntkctl_el1", POLICY_CLASS_TIMER, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, 0x3u, 1, 111 },
    { KEY(3, 3, 14, 3, 0), "cntv_tval_el0", POLICY_CLASS_TIMER, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, UINT64_MAX, 1, 112 },
    { KEY(3, 3, 14, 3, 1), "cntv_ctl_el0", POLICY_CLASS_TIMER, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, 0x7u, 1, 113 },
    { KEY(3, 3, 14, 3, 2), "cntv_cval_el0", POLICY_CLASS_TIMER, POLICY_EMULATE_READ, POLICY_EMULATE_WRITE, 0, UINT64_MAX, 1, 114 },
    { KEY(2, 0, 0, 2, 2), "mdscr_el1", POLICY_CLASS_DEBUG, POLICY_EMULATE_READ, POLICY_WRITE_MASK, 0, 0ull, 1, 115 },
};

static enum policy_profile g_profile;
static struct policy_entry g_effective_entry;

static bool policy_action_valid(enum policy_action action)
{
    return action == POLICY_ALLOW_NATIVE ||
           action == POLICY_EMULATE_READ ||
           action == POLICY_EMULATE_WRITE ||
           action == POLICY_WRITE_MASK ||
           action == POLICY_DENY_SKIP ||
           action == POLICY_DENY_UNDEF ||
           action == POLICY_PAUSE ||
           action == POLICY_PANIC;
}

static bool policy_class_valid(enum policy_access_class access_class)
{
    return access_class == POLICY_CLASS_ID ||
           access_class == POLICY_CLASS_MMU ||
           access_class == POLICY_CLASS_TIMER ||
           access_class == POLICY_CLASS_THREAD ||
           access_class == POLICY_CLASS_PMU ||
           access_class == POLICY_CLASS_DEBUG ||
           access_class == POLICY_CLASS_EL2_EL3 ||
           access_class == POLICY_CLASS_MISC;
}

static bool policy_lattice_valid(const struct policy_entry *entry)
{
    if (entry == (const struct policy_entry *)0 || entry->name == (const char *)0) {
        return false;
    }
    if (!policy_class_valid(entry->access_class) ||
        !policy_action_valid(entry->read_policy) ||
        !policy_action_valid(entry->write_policy)) {
        return false;
    }
    if (entry->access_class == POLICY_CLASS_EL2_EL3 &&
        (entry->read_policy != POLICY_DENY_UNDEF ||
         entry->write_policy != POLICY_DENY_UNDEF)) {
        return false;
    }
    if (entry->write_policy == POLICY_WRITE_MASK && entry->writable_mask == 0u &&
        entry->access_class != POLICY_CLASS_DEBUG) {
        return false;
    }
    return true;
}

static uint64_t policy_mask_value(const struct policy_entry *entry, uint64_t current, uint64_t operand)
{
    uint64_t preserve = current & ~entry->writable_mask;
    if (preserve == 0u) {
        preserve = entry->reset_value & ~entry->writable_mask;
    }
    return preserve | (operand & entry->writable_mask);
}

void policy_init(void)
{
    g_profile = POLICY_PROFILE_STRICT;
}

const char *policy_profile_name(enum policy_profile profile)
{
    switch (profile) {
    case POLICY_PROFILE_STRICT: return "strict";
    case POLICY_PROFILE_DIAGNOSTIC: return "diagnostic";
    case POLICY_PROFILE_LINUX_BOOT: return "linux-boot";
    case POLICY_PROFILE_DEBUG: return "debug";
    default: return "unknown";
    }
}

bool policy_set_profile(enum policy_profile profile)
{
    if (profile > POLICY_PROFILE_DEBUG) {
        return false;
    }
    g_profile = profile;
    return true;
}

bool policy_set_profile_name(const char *name)
{
    if (name == (const char *)0) {
        return false;
    }
    if (hv_strcmp(name, "strict") == 0) {
        return policy_set_profile(POLICY_PROFILE_STRICT);
    }
    if (hv_strcmp(name, "diagnostic") == 0) {
        return policy_set_profile(POLICY_PROFILE_DIAGNOSTIC);
    }
    if (hv_strcmp(name, "linux-boot") == 0) {
        return policy_set_profile(POLICY_PROFILE_LINUX_BOOT);
    }
    if (hv_strcmp(name, "debug") == 0) {
        return policy_set_profile(POLICY_PROFILE_DEBUG);
    }
    return false;
}

enum policy_profile policy_get_profile(void)
{
    return g_profile;
}

const struct policy_entry *sysreg_policy_lookup(uint64_t key)
{
    for (uint32_t i = 0; i < HV_ARRAY_SIZE(g_policy); i++) {
        if (g_policy[i].key == key) {
            return &g_policy[i];
        }
    }
    return (const struct policy_entry *)0;
}

static uint64_t *shadow_slot(struct vcpu *vcpu, uint64_t key)
{
    if (key == KEY(3, 0, 1, 0, 0)) return &vcpu->sysregs.sctlr_el1;
    if (key == KEY(3, 0, 1, 0, 2)) return &vcpu->sysregs.cpacr_el1;
    if (key == KEY(3, 0, 2, 0, 0)) return &vcpu->sysregs.ttbr0_el1;
    if (key == KEY(3, 0, 2, 0, 1)) return &vcpu->sysregs.ttbr1_el1;
    if (key == KEY(3, 0, 2, 0, 2)) return &vcpu->sysregs.tcr_el1;
    if (key == KEY(3, 0, 10, 2, 0)) return &vcpu->sysregs.mair_el1;
    if (key == KEY(3, 0, 12, 0, 0)) return &vcpu->sysregs.vbar_el1;
    if (key == KEY(3, 0, 13, 0, 1)) return &vcpu->sysregs.contextidr_el1;
    if (key == KEY(3, 0, 13, 0, 4)) return &vcpu->sysregs.tpidr_el1;
    if (key == KEY(3, 3, 13, 0, 3)) return &vcpu->sysregs.tpidrro_el0;
    if (key == KEY(3, 0, 14, 1, 0)) return &vcpu->sysregs.cntkctl_el1;
    if (key == KEY(3, 3, 14, 3, 0)) return &vcpu->sysregs.cntv_tval_el0;
    if (key == KEY(3, 3, 14, 3, 1)) return &vcpu->sysregs.cntv_ctl_el0;
    if (key == KEY(3, 3, 14, 3, 2)) return &vcpu->sysregs.cntv_cval_el0;
    if (key == KEY(3, 2, 0, 0, 0)) return &vcpu->sysregs.csselr_el1;
    if (key == KEY(2, 0, 0, 2, 2)) return &vcpu->sysregs.mdscr_el1;
    return (uint64_t *)0;
}

static bool is_el1_mmu_key(uint64_t key)
{
    return key == KEY(3, 0, 1, 0, 0) ||
           key == KEY(3, 0, 1, 0, 2) ||
           key == KEY(3, 0, 2, 0, 0) ||
           key == KEY(3, 0, 2, 0, 1) ||
           key == KEY(3, 0, 2, 0, 2) ||
           key == KEY(3, 0, 10, 2, 0) ||
           key == KEY(3, 0, 12, 0, 0) ||
           key == KEY(3, 0, 14, 1, 0);
}

static const struct policy_entry *effective_entry(const struct policy_entry *entry)
{
    if (entry == (const struct policy_entry *)0) {
        return (const struct policy_entry *)0;
    }

    g_effective_entry = *entry;
    if (g_profile == POLICY_PROFILE_DIAGNOSTIC) {
        if (g_effective_entry.access_class == POLICY_CLASS_PMU ||
            g_effective_entry.access_class == POLICY_CLASS_MISC) {
            g_effective_entry.read_policy = POLICY_DENY_SKIP;
            g_effective_entry.write_policy = POLICY_DENY_SKIP;
        }
    } else if (g_profile == POLICY_PROFILE_LINUX_BOOT) {
        if (g_effective_entry.access_class == POLICY_CLASS_PMU ||
            g_effective_entry.key == KEY(3, 0, 12, 0, 1)) {
            g_effective_entry.read_policy = POLICY_DENY_SKIP;
            g_effective_entry.write_policy = POLICY_DENY_SKIP;
        }
        if (g_effective_entry.access_class == POLICY_CLASS_MMU) {
            g_effective_entry.writable_mask = UINT64_MAX;
        }
        if (g_effective_entry.key == KEY(3, 3, 9, 12, 0)) {
            g_effective_entry.reset_value = 0u;
        }
    } else if (g_profile == POLICY_PROFILE_DEBUG) {
        if (g_effective_entry.access_class == POLICY_CLASS_DEBUG ||
            g_effective_entry.access_class == POLICY_CLASS_PMU) {
            g_effective_entry.read_policy = POLICY_EMULATE_READ;
            g_effective_entry.write_policy = POLICY_WRITE_MASK;
            g_effective_entry.writable_mask = 0u;
        }
    }
    return &g_effective_entry;
}

bool policy_apply_sysreg(struct vcpu *vcpu, const struct sysreg_access *access,
                         uint64_t operand, struct policy_decision *decision)
{
    if (vcpu == (struct vcpu *)0 || access == (const struct sysreg_access *)0 ||
        decision == (struct policy_decision *)0) {
        return false;
    }
    hv_memset(decision, 0, sizeof(*decision));
    const struct policy_entry *entry = effective_entry(sysreg_policy_lookup(access->key));
    decision->entry = entry;
    decision->action = POLICY_DENY_SKIP;
    decision->reason = 250u;

    if (entry == (const struct policy_entry *)0) {
        return true;
    }

    decision->action = access->is_read ? entry->read_policy : entry->write_policy;
    decision->reason = entry->violation_reason;

    uint64_t *slot = shadow_slot(vcpu, access->key);
    switch (decision->action) {
    case POLICY_EMULATE_READ:
        decision->value = slot != (uint64_t *)0 ? *slot : entry->reset_value;
        return true;
    case POLICY_EMULATE_WRITE:
        if (slot == (uint64_t *)0) {
            decision->action = POLICY_DENY_UNDEF;
            return true;
        }
        *slot = policy_mask_value(entry, *slot, operand);
        decision->value = *slot;
        if (g_profile == POLICY_PROFILE_LINUX_BOOT && is_el1_mmu_key(access->key)) {
            arch_sync_el1_sysregs(vcpu);
        }
        timer_vcpu_sync(vcpu);
        return true;
    case POLICY_WRITE_MASK:
        if (slot == (uint64_t *)0) {
            decision->action = POLICY_DENY_UNDEF;
            return true;
        }
        *slot = policy_mask_value(entry, *slot, operand);
        decision->value = *slot;
        if (g_profile == POLICY_PROFILE_LINUX_BOOT && is_el1_mmu_key(access->key)) {
            arch_sync_el1_sysregs(vcpu);
        }
        timer_vcpu_sync(vcpu);
        return true;
    case POLICY_DENY_SKIP:
    case POLICY_DENY_UNDEF:
    case POLICY_PAUSE:
    case POLICY_PANIC:
    case POLICY_ALLOW_NATIVE:
    default:
        return true;
    }
}

void policy_dump(void)
{
    uart_puts("policy profile=");
    uart_puts(policy_profile_name(g_profile));
    uart_puts("\n");
    for (uint32_t i = 0; i < HV_ARRAY_SIZE(g_policy); i++) {
        const struct policy_entry *entry = effective_entry(&g_policy[i]);
        uart_puts("policy key=");
        uart_put_hex64(entry->key);
        uart_puts(" name=");
        uart_puts(entry->name);
        uart_puts(" read=");
        uart_put_dec64((uint64_t)entry->read_policy);
        uart_puts(" write=");
        uart_put_dec64((uint64_t)entry->write_policy);
        uart_puts(" mask=");
        uart_put_hex64(entry->writable_mask);
        uart_puts("\n");
    }
}

bool policy_selftest(void)
{
    enum policy_profile saved = g_profile;
    bool profile_ok;
    for (uint32_t i = 0; i < HV_ARRAY_SIZE(g_policy); i++) {
        if (!policy_lattice_valid(&g_policy[i])) {
            return false;
        }
        if (sysreg_policy_lookup(g_policy[i].key) != &g_policy[i]) {
            return false;
        }
        for (uint32_t j = i + 1u; j < HV_ARRAY_SIZE(g_policy); j++) {
            if (g_policy[i].key == g_policy[j].key) {
                return false;
            }
        }
    }
    profile_ok = policy_set_profile_name("linux-boot") &&
                 policy_get_profile() == POLICY_PROFILE_LINUX_BOOT &&
                 policy_set_profile(saved);
    return sysreg_policy_lookup(KEY(3, 4, 1, 1, 0)) != (const struct policy_entry *)0 &&
           sysreg_policy_lookup(KEY(3, 7, 15, 15, 7)) == (const struct policy_entry *)0 &&
           profile_ok;
}