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

_Static_assert(offsetof(struct trap_frame, x) == 0u, "trap_frame x offset");
_Static_assert(offsetof(struct trap_frame, sp_el0) == 248u, "trap_frame sp_el0 offset");
_Static_assert(offsetof(struct trap_frame, sp_el1) == 256u, "trap_frame sp_el1 offset");
_Static_assert(offsetof(struct trap_frame, elr_el2) == 264u, "trap_frame elr offset");
_Static_assert(offsetof(struct trap_frame, spsr_el2) == 272u, "trap_frame spsr offset");
_Static_assert(offsetof(struct trap_frame, esr_el2) == 280u, "trap_frame esr offset");
_Static_assert(offsetof(struct trap_frame, far_el2) == 288u, "trap_frame far offset");
_Static_assert(offsetof(struct trap_frame, hpfar_el2) == 296u, "trap_frame hpfar offset");
_Static_assert(sizeof(struct trap_frame) == 304u, "trap_frame size");

_Static_assert(offsetof(struct vcpu, x) == 0u, "vcpu x offset");
_Static_assert(offsetof(struct vcpu, sp_el0) == 248u, "vcpu sp_el0 offset");
_Static_assert(offsetof(struct vcpu, sp_el1) == 256u, "vcpu sp_el1 offset");
_Static_assert(offsetof(struct vcpu, pc) == 264u, "vcpu pc offset");
_Static_assert(offsetof(struct vcpu, pstate) == 272u, "vcpu pstate offset");

static struct vcpu *g_current_vcpu;

void vcpu_init(struct vcpu *vcpu, uint32_t id, uint64_t entry, uint64_t stack_top)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    hv_memset(vcpu, 0, sizeof(*vcpu));
    vcpu->id = id;
    vcpu->state = VCPU_CREATED;
    vcpu->pc = entry;
    vcpu->pstate = 0x3c5u;
    vcpu->sp_el1 = stack_top;
    vcpu->sysregs.sctlr_el1 = 0x30d00800ull;
    vcpu->sysregs.mair_el1 = 0xffull;
    vcpu->sysregs.cntkctl_el1 = 0u;
    vcpu->pending_virtual_irq = 0u;
    vcpu->active_virtual_irq = 0u;
    vcpu->state = VCPU_READY;
}

void vcpu_load_frame(struct vcpu *vcpu, const struct trap_frame *frame)
{
    BUG_ON(vcpu == (struct vcpu *)0 || frame == (const struct trap_frame *)0);
    for (uint32_t i = 0; i < 31u; i++) {
        vcpu->x[i] = frame->x[i];
    }
    vcpu->sp_el0 = frame->sp_el0;
    vcpu->sp_el1 = frame->sp_el1;
    vcpu->pc = frame->elr_el2;
    vcpu->pstate = frame->spsr_el2;
    vcpu->last_esr = frame->esr_el2;
    vcpu->last_far = frame->far_el2;
    vcpu->last_hpfar = frame->hpfar_el2;
    arch_capture_el1_sysregs(vcpu);
}

void vcpu_store_frame(const struct vcpu *vcpu, struct trap_frame *frame)
{
    BUG_ON(vcpu == (const struct vcpu *)0 || frame == (struct trap_frame *)0);
    for (uint32_t i = 0; i < 31u; i++) {
        frame->x[i] = vcpu->x[i];
    }
    frame->sp_el0 = vcpu->sp_el0;
    frame->sp_el1 = vcpu->sp_el1;
    frame->elr_el2 = vcpu->pc;
    frame->spsr_el2 = vcpu->pstate;
}

uint64_t vcpu_read_gpr(const struct vcpu *vcpu, uint32_t rt)
{
    BUG_ON(vcpu == (const struct vcpu *)0);
    if (rt >= 31u) {
        return 0u;
    }
    return vcpu->x[rt];
}

void vcpu_write_gpr(struct vcpu *vcpu, uint32_t rt, uint64_t value)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    if (rt < 31u) {
        vcpu->x[rt] = value;
    }
}

void vcpu_set_pending_irq(struct vcpu *vcpu, uint32_t intid)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    if (intid < 64u) {
        vcpu->pending_virtual_irq |= (1ull << intid);
    }
}

void vcpu_clear_pending_irq(struct vcpu *vcpu, uint32_t intid)
{
    BUG_ON(vcpu == (struct vcpu *)0);
    if (intid < 64u) {
        vcpu->pending_virtual_irq &= ~(1ull << intid);
    }
}

struct vcpu *vcpu_current(void)
{
    return g_current_vcpu;
}

void vcpu_set_current(struct vcpu *vcpu)
{
    g_current_vcpu = vcpu;
}

static void vcpu_dump_exception_telemetry(const struct vcpu *vcpu)
{
    uart_puts("trap_ec ec=");
    uart_put_hex64(ESR_EC_HVC64);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_HVC64));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.hvc_calls);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_SYSREG);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_SYSREG));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.sysreg_traps);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_SMC64);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_SMC64));
    uart_puts(" denied=");
    uart_put_dec64(vcpu->stats.smc_denied);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_WFI_WFE);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_WFI_WFE));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.wfx_traps);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_IABT_LOWER);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_IABT_LOWER));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.instruction_aborts);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_DABT_LOWER);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_DABT_LOWER));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.data_aborts);
    uart_puts("\ntrap_ec ec=");
    uart_put_hex64(ESR_EC_UNKNOWN);
    uart_puts(" name=");
    uart_puts(esr_ec_name(ESR_EC_UNKNOWN));
    uart_puts(" count=");
    uart_put_dec64(vcpu->stats.unknown_exceptions);
    uart_puts("\n");
}

void vcpu_dump(const struct vcpu *vcpu)
{
    if (vcpu == (const struct vcpu *)0) {
        uart_puts("vcpu none\n");
        return;
    }
    uart_puts("vcpu id=");
    uart_put_dec64(vcpu->id);
    uart_puts(" state=");
    uart_put_dec64((uint64_t)vcpu->state);
    uart_puts(" pc=");
    uart_put_hex64(vcpu->pc);
    uart_puts(" pstate=");
    uart_put_hex64(vcpu->pstate);
    uart_puts(" sp_el1=");
    uart_put_hex64(vcpu->sp_el1);
    uart_puts(" traps=");
    uart_put_dec64(vcpu->stats.traps);
    uart_puts(" last_ec=");
    uart_put_hex64(esr_ec(vcpu->last_esr));
    uart_puts(" last_ecname=");
    uart_puts(esr_ec_name(esr_ec(vcpu->last_esr)));
    uart_puts(" injected_undefs=");
    uart_put_dec64(vcpu->stats.injected_undefs);
    uart_puts(" pending_irq=");
    uart_put_hex64(vcpu->pending_virtual_irq);
    uart_puts(" virq_inj=");
    uart_put_dec64(vcpu->stats.virtual_irqs_injected);
    uart_puts(" virq_done=");
    uart_put_dec64(vcpu->stats.virtual_irqs_completed);
    uart_puts(" timer_syncs=");
    uart_put_dec64(vcpu->stats.timer_syncs);
    uart_puts(" timer_delta=");
    uart_put_dec64(vcpu->stats.timer_last_delta);
    uart_puts("\n");
    vcpu_dump_exception_telemetry(vcpu);
    for (uint32_t i = 0; i < 31u; i++) {
        uart_puts("x");
        uart_put_dec64(i);
        uart_puts("=");
        uart_put_hex64(vcpu->x[i]);
        uart_puts((i % 4u) == 3u ? "\n" : " ");
    }
    uart_puts("\n");
}

bool vcpu_frame_selftest(void)
{
    struct trap_frame frame;
    struct trap_frame out;
    struct vcpu vcpu;

    hv_memset(&frame, 0, sizeof(frame));
    hv_memset(&out, 0, sizeof(out));
    hv_memset(&vcpu, 0, sizeof(vcpu));
    for (uint32_t i = 0; i < 31u; i++) {
        frame.x[i] = 0x1000u + i;
    }
    frame.sp_el0 = 0x2000u;
    frame.sp_el1 = 0x3000u;
    frame.elr_el2 = 0x4000u;
    frame.spsr_el2 = 0x3c5u;
    frame.esr_el2 = 0x96000000u;
    frame.far_el2 = 0x5000u;
    frame.hpfar_el2 = 0x6000u;

    vcpu_load_frame(&vcpu, &frame);
    vcpu_store_frame(&vcpu, &out);
    return hv_memcmp(frame.x, out.x, sizeof(frame.x)) == 0 &&
           frame.sp_el0 == out.sp_el0 &&
           frame.sp_el1 == out.sp_el1 &&
           frame.elr_el2 == out.elr_el2 &&
           frame.spsr_el2 == out.spsr_el2 &&
           vcpu.last_esr == frame.esr_el2 &&
           vcpu.last_far == frame.far_el2 &&
           vcpu.last_hpfar == frame.hpfar_el2;
}