/* hypercall.c — UTSM-side VMCALL hypercall handler.
 *
 * When the Linux guest executes a VMCALL instruction, it triggers a VM-Exit
 * with exit reason EXIT_VCALL (18). The guest places:
 *   RAX = UTSM_HCALL_MAGIC | op    (identifies this as a UTSM hypercall)
 *   RDI = arg0
 *   RSI = arg1
 *   RDX = arg2
 *
 * The VM-Exit assembly stub saves all guest GPRs into g_guest_regs before
 * calling the C dispatcher. This handler reads arguments from g_guest_regs,
 * dispatches to the operation handler, writes the return value into
 * g_guest_regs.rax, and advances GUEST_RIP so the guest resumes after
 * the VMCALL instruction.
 *
 * Protocol details in CODE/utsm-ipc/ipc_proto.h.
 */

#include <utsm/hypercall.h>
#include <utsm/vmx.h>
#include <utsm/ept.h>
#include <utsm/vmm.h>
#include <utsm/ipc_shm.h>
#include <utsm/log.h>
#include <utsm/types.h>
#include <ipc_proto.h>
#include "../arch/x86_64/limine.h"

extern volatile struct limine_hhdm_request g_hhdm_request;

static u64 g_hypercall_count;

/* ===== Guest memory access helpers =====
 *
 * The Linux guest uses identity-mapped page tables (GVA == GPA), so
 * pointers passed via hypercall arguments are GPAs. We translate GPA→HPA
 * via EPT walk, then add HHDM offset to get the host virtual address. */

static void *gpa_to_hva(u64 gpa) {
    u64 hpa = ept_gpa_to_hpa(gpa);
    if (hpa == 0) return (void *)0;
    if (!g_hhdm_request.response) return (void *)0;
    u64 hhdm = g_hhdm_request.response->offset;
    return (void *)(hhdm + hpa);
}

/* Write a u64 to guest memory at the given GPA.
 * Returns 0 on success, -1 on failure. */
static int guest_write_u64(u64 gpa, u64 value) {
    void *hva = gpa_to_hva(gpa);
    if (!hva) return -1;
    *(volatile u64 *)hva = value;
    return 0;
}

/* Copy bytes from guest memory (GVA=GPA) to host buffer.
 * Handles cross-page boundaries by translating each page separately.
 * Returns 0 on success, -1 on failure. */
static int guest_read_buf(u64 gpa, void *host_dst, u64 len) {
    u8 *dst = (u8 *)host_dst;
    u64 copied = 0;
    while (copied < len) {
        void *hva = gpa_to_hva(gpa + copied);
        if (!hva) return -1;
        u64 page_end = (gpa + copied + 4096) & ~0xFFFULL;
        u64 chunk = page_end - (gpa + copied);
        if (chunk > len - copied) chunk = len - copied;
        const u8 *src = (const u8 *)hva;
        for (u64 i = 0; i < chunk; i++) {
            dst[copied + i] = src[i];
        }
        copied += chunk;
    }
    return 0;
}

/* ===== Individual hypercall handlers ===== */

static long hcall_ping(u64 a0, u64 a1, u64 a2) {
    (void)a0; (void)a1; (void)a2;
    /* Return a recognizable PONG value */
    return 0x504F4E47;  /* "PONG" */
}

static long hcall_hello(u64 a0, u64 a1, u64 a2) {
    (void)a0; (void)a1; (void)a2;
    log_info("[HCALL] HELLO from Linux guest");

    /* Linux is announcing itself; the IPC shared memory's linux_ready
     * flag is set separately when Linux maps the shared memory. */
    if (ipc_shm_is_ready()) {
        /* Send a HELLO_ACK message via the shared memory ring */
        const char *msg = "HELLO from UTSM";
        ipc_shm_send(UTSM_MSG_HELLO_ACK, msg, 15);
    }
    return UTSM_HCALL_OK;
}

static long hcall_shm_info(u64 a0, u64 a1, u64 a2) {
    /* a0 = GPA of gpa_out (u64), a1 = GPA of size_out (u64) */
    (void)a2;
    if (!ipc_shm_is_ready()) {
        return UTSM_HCALL_NOMEM;
    }

    u64 shm_gpa = ipc_shm_get_gpa();
    u64 shm_size = ipc_shm_get_size();

    /* Write results to guest memory */
    if (a0 != 0 && guest_write_u64(a0, shm_gpa) != 0) {
        log_error("[HCALL] SHM_INFO: failed to write gpa to guest");
        return UTSM_HCALL_INVAL;
    }
    if (a1 != 0 && guest_write_u64(a1, shm_size) != 0) {
        log_error("[HCALL] SHM_INFO: failed to write size to guest");
        return UTSM_HCALL_INVAL;
    }

    log_hex64("[HCALL] SHM_INFO: gpa=", shm_gpa);
    log_hex64("[HCALL] SHM_INFO: size=", shm_size);
    return UTSM_HCALL_OK;
}

static long hcall_cap_validate(u64 a0, u64 a1, u64 a2) {
    /* Phase 1.3: stub — capability validation not yet implemented */
    (void)a0; (void)a1; (void)a2;
    log_info("[HCALL] CAP_VALIDATE (stub)");
    return UTSM_HCALL_NOSYS;
}

static long hcall_utrw_read(u64 a0, u64 a1, u64 a2) {
    /* Phase 1.3: stub — sealed memory read not yet implemented */
    (void)a0; (void)a1; (void)a2;
    log_info("[HCALL] UTRW_READ (stub)");
    return UTSM_HCALL_NOSYS;
}

static long hcall_utrw_write(u64 a0, u64 a1, u64 a2) {
    /* Phase 1.3: stub — sealed memory write not yet implemented */
    (void)a0; (void)a1; (void)a2;
    log_info("[HCALL] UTRW_WRITE (stub)");
    return UTSM_HCALL_NOSYS;
}

static long hcall_drr_checkpoint(u64 a0, u64 a1, u64 a2) {
    /* Phase 1.3: stub — DRR checkpoint trigger not yet implemented */
    (void)a0; (void)a1; (void)a2;
    log_info("[HCALL] DRR_CHECKPOINT (stub)");
    return UTSM_HCALL_NOSYS;
}

static long hcall_console_write(u64 a0, u64 a1, u64 a2) {
    /* a0 = GPA of string buffer, a1 = length */
    (void)a2;
    if (a0 == 0 || a1 == 0) {
        return UTSM_HCALL_INVAL;
    }

    /* Cap length to prevent runaway reads */
    if (a1 > 4096) a1 = 4096;

    /* Read the string from guest memory */
    char buf[4097];
    if (guest_read_buf(a0, buf, a1) != 0) {
        log_error("[HCALL] CONSOLE_WRITE: failed to read guest buffer");
        return UTSM_HCALL_INVAL;
    }
    buf[a1] = '\0';

    /* Output to UTSM serial console with a [LINUX] prefix */
    serial_write("[LINUX] ");
    serial_write(buf);
    return (long)a1;
}

static long hcall_console_read(u64 a0, u64 a1, u64 a2) {
    /* Phase 1.3: stub — console read not yet implemented */
    (void)a0; (void)a1; (void)a2;
    return UTSM_HCALL_NOSYS;
}

/* ===== Main hypercall dispatcher ===== */

int hypercall_handle(u64 guest_rax, u64 guest_rdi, u64 guest_rsi, u64 guest_rdx,
                     u64 guest_rip, u64 instr_len) {
    g_hypercall_count++;

    /* Validate magic: the high bits of RAX must be UTSM_HCALL_MAGIC */
    if ((guest_rax & ~0xFFFFULL) != UTSM_HCALL_MAGIC) {
        log_hex64("[HCALL] bad magic: ", guest_rax);
        /* Return error in guest RAX */
        g_guest_regs.rax = (u64)UTSM_HCALL_INVAL;
        /* Still advance RIP to avoid infinite loop */
        vmx_vmcs_write(VMCS_GUEST_RIP, guest_rip + instr_len);
        return 1;  /* resume guest */
    }

    /* Extract operation from low 16 bits */
    u64 op = guest_rax & 0xFFFFULL;
    long result = UTSM_HCALL_NOSYS;

    switch (op) {
    case UTSM_HCALL_PING:
        result = hcall_ping(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_HELLO:
        result = hcall_hello(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_SHM_INFO:
        result = hcall_shm_info(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_CAP_VALIDATE:
        result = hcall_cap_validate(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_UTRW_READ:
        result = hcall_utrw_read(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_UTRW_WRITE:
        result = hcall_utrw_write(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_DRR_CHECKPOINT:
        result = hcall_drr_checkpoint(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_CONSOLE_WRITE:
        result = hcall_console_write(guest_rdi, guest_rsi, guest_rdx);
        break;
    case UTSM_HCALL_CONSOLE_READ:
        result = hcall_console_read(guest_rdi, guest_rsi, guest_rdx);
        break;
    default:
        log_hex64("[HCALL] unknown op: ", op);
        result = UTSM_HCALL_NOSYS;
        break;
    }

    /* Write return value into guest RAX (restored by assembly before vmresume) */
    g_guest_regs.rax = (u64)result;

    /* Advance guest RIP past the VMCALL instruction */
    vmx_vmcs_write(VMCS_GUEST_RIP, guest_rip + instr_len);

    return 1;  /* resume guest */
}

u64 hypercall_get_count(void) {
    return g_hypercall_count;
}