/*
  StarForth — Steady-State Virtual Machine Runtime

  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.

  You may obtain a copy of the License at:
      https://github.com/star.4th@proton.me/StarForth/LICENSE.txt

  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose, and noninfringement.

 See the License for the specific language governing permissions and
 limitations under the License.

 StarForth — Steady-State Virtual Machine Runtime
  Copyright (c) 2023–2025 Robert A. James
  All rights reserved.

  This file is part of the StarForth project.

  Licensed under the StarForth License, Version 1.0 (the "License");
  you may not use this file except in compliance with the License.

  You may obtain a copy of the License at:
       https://github.com/star.4th@proton.me/StarForth/LICENSE.txt

  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose, and noninfringement.

  See the License for the specific language governing permissions and
  limitations under the License.

 */

/**
 * arena.c - StarKernel VM arena initialization
 *
 * Provides kernel-specific VM initialization with PMM-backed arena.
 * This replaces the hosted vm_init() for kernel builds.
 *
 * M7 Non-Negotiable: VM arena is PMM-backed contiguous pages.
 */

#ifdef __STARKERNEL__

#include "starkernel/vm/arena.h"
#include "vm_host.h"
#include "pmm.h"
#include "vmm.h"
#include "console.h"
#include "starkernel/hal/hal.h"
#include "dictionary_management.h"
#include <stddef.h>
#include <stdint.h>

/* Forward declaration - we'll need the VM struct */
/* For now, define minimal interface; full integration comes later */

/* aarch64/riscv64: UEFI provides 1:1 phys=virt mapping after ExitBootServices.
 * VMM installs page tables only on amd64 (via CR3); on aarch64/riscv64 the
 * high virtual address is never mapped.  Use physical address directly. */
#if defined(ARCH_AARCH64) || defined(__riscv)
#define SK_VM_ARENA_VADDR        0x0ULL  /* sentinel: use paddr as vaddr */
#else
#define SK_VM_ARENA_VADDR        0xFFFF900000000000ULL
#endif
#define SK_VM_GUARD_SIZE         PMM_PAGE_SIZE
#define SK_VM_GUARD_PATTERN_HEAD 0x5ac0ffeed0c0ffeeULL
#define SK_VM_GUARD_PATTERN_TAIL 0x0bad0badf00df00dULL

/* VM memory size from hosted VM (5 MB) */
#define SK_VM_MEMORY_SIZE   (5ULL * 1024ULL * 1024ULL)
#define SK_VM_TOTAL_SIZE    (SK_VM_MEMORY_SIZE + (2ULL * SK_VM_GUARD_SIZE))
#define SK_VM_ARENA_PAGES   ((SK_VM_TOTAL_SIZE + PMM_PAGE_SIZE - 1) / PMM_PAGE_SIZE)

/* Arena state */
static uint64_t vm_arena_paddr = 0;
static uint64_t vm_arena_guard_vaddr = 0;
static uint64_t vm_arena_client_vaddr = 0;
static int vm_arena_initialized = 0;

static void guard_fill(uint64_t base)
{
    uint64_t *ptr = (uint64_t *)(uintptr_t)base;
    ptr[0] = SK_VM_GUARD_PATTERN_HEAD;
    ptr[(SK_VM_GUARD_SIZE / sizeof(uint64_t)) - 1] = SK_VM_GUARD_PATTERN_TAIL;
}

static int guard_check(uint64_t base)
{
    const uint64_t *ptr = (const uint64_t *)(uintptr_t)base;
    if (ptr[0] != SK_VM_GUARD_PATTERN_HEAD) {
        return 0;
    }
    if (ptr[(SK_VM_GUARD_SIZE / sizeof(uint64_t)) - 1] != SK_VM_GUARD_PATTERN_TAIL) {
        return 0;
    }
    return 1;
}

/**
 * sk_vm_arena_alloc - Allocate PMM-backed VM arena
 *
 * Allocates contiguous physical pages and maps them RW+NX.
 * Returns virtual address of arena, or 0 on failure.
 */
uint64_t sk_vm_arena_alloc(void) {
    uint64_t paddr, guard_vaddr, client_vaddr;

    /* Fresh PMM allocation every call — each VM owns its own arena.
     * The singleton guard was removed because baby VMs (Artemis, Hermes)
     * share Mama's physical pages when sk_vm_arena_alloc returns the cached
     * address; their vm->here=0 at init overwrites Mama's compiled words and
     * their FILL operations zero Mama's word bodies (riscv64 K-PUSH crash). */
    paddr = pmm_alloc_contiguous(SK_VM_ARENA_PAGES);
    if (paddr == 0) {
        console_println("sk_vm_arena_alloc: PMM allocation failed");
        return 0;
    }

#if defined(__aarch64__) || defined(_M_ARM64) || defined(ARCH_AARCH64) || defined(__riscv)
    /* On aarch64/riscv64, UEFI's identity-mapped page tables remain active
     * after ExitBootServices(); physical address == virtual address for all RAM.
     * Each pmm_alloc_contiguous call returns a distinct paddr, so each VM gets
     * a distinct vaddr automatically. */
    guard_vaddr  = paddr;
    client_vaddr = guard_vaddr + SK_VM_GUARD_SIZE;
    (void)(VMM_FLAG_WRITABLE); (void)(VMM_FLAG_NX);
#else
    /* On amd64, walk a per-call virtual address pointer forward so each arena
     * gets a unique VMM mapping rather than re-mapping the same fixed address. */
    static uint64_t next_arena_vaddr = SK_VM_ARENA_VADDR;
    guard_vaddr  = next_arena_vaddr - SK_VM_GUARD_SIZE;
    client_vaddr = guard_vaddr + SK_VM_GUARD_SIZE;
    next_arena_vaddr += SK_VM_TOTAL_SIZE;
    uint64_t flags = VMM_FLAG_WRITABLE | VMM_FLAG_NX;

    if (vmm_map_range(guard_vaddr, paddr,
                      SK_VM_ARENA_PAGES * PMM_PAGE_SIZE, flags) != 0) {
        console_println("sk_vm_arena_alloc: VMM mapping failed");
        pmm_free_contiguous(paddr, SK_VM_ARENA_PAGES);
        next_arena_vaddr -= SK_VM_TOTAL_SIZE;
        return 0;
    }
#endif

    /* Zero the arena */
    uint8_t *arena = (uint8_t *)(uintptr_t)client_vaddr;
    for (size_t i = 0; i < SK_VM_MEMORY_SIZE; i++) {
        arena[i] = 0;
    }

    guard_fill(guard_vaddr);
    guard_fill(guard_vaddr + SK_VM_GUARD_SIZE + SK_VM_MEMORY_SIZE);

    /* Track first (Mama) arena for sk_vm_arena_free / sk_vm_arena_ptr / assert_guards */
    if (!vm_arena_initialized) {
        vm_arena_paddr        = paddr;
        vm_arena_guard_vaddr  = guard_vaddr;
        vm_arena_client_vaddr = client_vaddr;
        vm_arena_initialized  = 1;
    }

    console_puts("VM arena allocated: ");
    /* Print address (simple hex) */
    {
        char buf[20];
        uint64_t v = client_vaddr;
        int i = 18;
        buf[19] = '\0';
        buf[18] = '\n';
        while (i > 2) {
            int d = v & 0xF;
            buf[--i] = (d < 10) ? ('0' + d) : ('a' + d - 10);
            v >>= 4;
        }
        buf[1] = 'x';
        buf[0] = '0';
        console_puts(buf);
    }
    console_puts(" (");
    /* Print size in MB */
    {
        char buf[16];
        uint64_t mb = SK_VM_MEMORY_SIZE / (1024 * 1024);
        int i = 0;
        if (mb == 0) {
            buf[i++] = '0';
        } else {
            char tmp[16];
            int j = 0;
            while (mb > 0) {
                tmp[j++] = '0' + (mb % 10);
                mb /= 10;
            }
            while (j > 0) {
                buf[i++] = tmp[--j];
            }
        }
        buf[i] = '\0';
        console_puts(buf);
    }
    console_println(" MB)");

    return client_vaddr;
}

/**
 * sk_vm_arena_free - Free the VM arena
 *
 * Unmaps and frees the PMM-backed arena.
 */
void sk_vm_arena_free(void) {
    if (!vm_arena_initialized) {
        return;
    }

#if !(defined(__aarch64__) || defined(_M_ARM64) || defined(ARCH_AARCH64) || defined(__riscv))
    /* Unmap each page (not needed on aarch64/riscv64 — no VMM remapping was done) */
    for (uint64_t offset = 0; offset < SK_VM_TOTAL_SIZE; offset += PMM_PAGE_SIZE) {
        vmm_unmap_page(vm_arena_guard_vaddr + offset);
    }
#endif

    /* Free physical memory */
    pmm_free_contiguous(vm_arena_paddr, SK_VM_ARENA_PAGES);

    vm_arena_paddr = 0;
    vm_arena_guard_vaddr = 0;
    vm_arena_client_vaddr = 0;
    vm_arena_initialized = 0;
}

/**
 * sk_vm_arena_ptr - Get pointer to VM arena
 *
 * Returns NULL if arena not allocated.
 */
uint8_t* sk_vm_arena_ptr(void) {
    if (!vm_arena_initialized) {
        return (uint8_t*)0;
    }
    return (uint8_t *)(uintptr_t)vm_arena_client_vaddr;
}

/**
 * sk_vm_arena_size - Get VM arena size
 */
size_t sk_vm_arena_size(void) {
    return SK_VM_MEMORY_SIZE;
}

/**
 * sk_vm_arena_is_initialized - Check if arena is ready
 */
int sk_vm_arena_is_initialized(void) {
    return vm_arena_initialized;
}

/* ===========================================================================
 * Per-VM native stack (FABRIC-3.md §XXVIII, Stage 1, 2026-09-13)
 * ===========================================================================
 *
 * 2 MiB, matching this project's own established kernel-stack convention
 * (g_kernel_stack / g_rpi5_native_stack in each arch's entry assembly) --
 * not a guessed number. That single 2 MiB stack today accommodates the
 * ENTIRE shared kernel C stack's worth of nested VM-EXEC/execute_colon_word
 * recursion across every VM combined, so 2 MiB for one VM's own stack alone
 * is generous headroom, not a guess at the edge of what's needed.
 */
#define SK_VM_STACK_SIZE   (2ULL * 1024ULL * 1024ULL)
#define SK_VM_STACK_PAGES  ((SK_VM_STACK_SIZE + (2ULL * SK_VM_GUARD_SIZE) + PMM_PAGE_SIZE - 1) / PMM_PAGE_SIZE)

#if !(defined(ARCH_AARCH64) || defined(__riscv))
/* amd64 needs its own virtual-address walk, separate from arena's
 * next_arena_vaddr -- different logical region, would otherwise collide. */
#define SK_VM_STACK_ARENA_VADDR  0xFFFF980000000000ULL
static uint64_t next_stack_vaddr = SK_VM_STACK_ARENA_VADDR;
#endif

/**
 * sk_vm_native_stack_alloc - allocate one VM's own dedicated native stack.
 *
 * Always a fresh, independent PMM allocation with real guard pages -- no
 * singleton, no kmalloc fallback, unlike sk_vm_arena_alloc()'s Mama-only
 * special case. Returns 0 on success (fields of *out populated), -1 on
 * failure.
 */
int sk_vm_native_stack_alloc(sk_vm_native_stack_t *out) {
    uint64_t paddr, guard_vaddr;

    if (!out) return -1;

    paddr = pmm_alloc_contiguous(SK_VM_STACK_PAGES);
    if (paddr == 0) {
        console_println("sk_vm_native_stack_alloc: PMM allocation failed");
        return -1;
    }

#if defined(ARCH_AARCH64) || defined(__riscv)
    guard_vaddr = paddr;
    (void)(VMM_FLAG_WRITABLE); (void)(VMM_FLAG_NX);
#else
    guard_vaddr = next_stack_vaddr;
    next_stack_vaddr += SK_VM_STACK_PAGES * PMM_PAGE_SIZE;
    uint64_t flags = VMM_FLAG_WRITABLE | VMM_FLAG_NX;

    if (vmm_map_range(guard_vaddr, paddr,
                      SK_VM_STACK_PAGES * PMM_PAGE_SIZE, flags) != 0) {
        console_println("sk_vm_native_stack_alloc: VMM mapping failed");
        pmm_free_contiguous(paddr, SK_VM_STACK_PAGES);
        next_stack_vaddr -= SK_VM_STACK_PAGES * PMM_PAGE_SIZE;
        return -1;
    }
#endif

    guard_fill(guard_vaddr);
    guard_fill(guard_vaddr + SK_VM_GUARD_SIZE + SK_VM_STACK_SIZE);

    out->paddr       = paddr;
    out->guard_vaddr = guard_vaddr;
    out->stack_top   = guard_vaddr + SK_VM_GUARD_SIZE + SK_VM_STACK_SIZE;
    return 0;
}

/**
 * sk_vm_native_stack_free - release a VM's native stack.
 *
 * Mirrors sk_vm_arena_free()'s unmap/free sequence, but per-call (this is
 * not a singleton) -- safe to call once per successful
 * sk_vm_native_stack_alloc(), and the caller (vm_cleanup()) is responsible
 * for not calling it twice on the same struct.
 */
void sk_vm_native_stack_free(const sk_vm_native_stack_t *stack) {
    if (!stack || stack->paddr == 0) return;

#if !(defined(ARCH_AARCH64) || defined(__riscv))
    for (uint64_t offset = 0; offset < SK_VM_STACK_PAGES * PMM_PAGE_SIZE; offset += PMM_PAGE_SIZE) {
        vmm_unmap_page(stack->guard_vaddr + offset);
    }
#endif

    pmm_free_contiguous(stack->paddr, SK_VM_STACK_PAGES);
}

void sk_vm_arena_assert_guards(const char *tag) {
    if (!vm_arena_initialized) {
        return;
    }
    if (!guard_check(vm_arena_guard_vaddr) ||
        !guard_check(vm_arena_guard_vaddr + SK_VM_GUARD_SIZE + SK_VM_MEMORY_SIZE)) {
        console_puts("[VM][arena] guard corruption at ");
        console_println(tag ? tag : "unknown");
        /* vm_dictionary_log_last_word(NULL, tag); -- not implemented in minimal kernel */
        sk_hal_panic("VM arena guard corrupted");
    }
}

#endif /* __STARKERNEL__ */