/*
  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.
*/

/**
 * capsule_loader.c - Block Capsule Loader and Executor (M7.1)
 *
 * Parses a capsule payload for "Block <num>" headers and implements
 * the full init sequence:
 *
 *   1. Write blocks to ramdrive (2048+) — establishes device content
 *      so FORTH can call LOAD N later if needed.
 *   2. Execute each block directly from payload bytes via exec_block_with_retry.
 *      exec_block_with_retry receives a C pointer into the payload; it does
 *      NOT use the block device, so no dedicated-RAM copy is needed here.
 *   3. Zero dedicated RAM slots so userspace can reuse them.
 *
 * The ramdrive (2048+) is the fallback device when no physical block
 * device is mounted.
 */

#include "starkernel/capsule_loader.h"
#include "starkernel/capsule.h"
#include "starkernel/capsule_run.h"
#include "block_subsystem.h"
#include "vm.h"
#ifdef __STARKERNEL__
#include "starkernel/console.h"
#include "log.h"
#endif

/* Ramdrive starts at this block number; dest = source - CAPSULE_RAM_OFFSET */
#define CAPSULE_RAM_OFFSET   2048u
#define LOADER_BLOCK_SIZE    1024u

/* FORTH block geometry: 16 lines × 64 characters = 1024 bytes */
#define CAPSULE_BLOCK_LINES  16
#define CAPSULE_LINE_COLS    64

/*
 * Execution line buffer — capsule text (the .4th source representation)
 * may use lines longer than the canonical 64-column block storage limit.
 * This buffer covers the execution path; CAPSULE_LINE_COLS covers storage
 * and the 16-bit bitmask geometry.
 */
#define CAPSULE_EXEC_LINE_MAX 256

/*
 * Forward-definition retry budget.
 *
 * When a block line triggers an error (typically UNKNOWN WORD for a
 * forward reference) the loader skips that line and retries the entire
 * block from the top.  Each retry sets one bit in a uint16_t bitmask
 * (one bit per line 0-15).  When the mask reaches 0xFFFF all 16 line
 * slots are exhausted and the block is declared unloadable; the loader
 * panics with a per-line report.
 *
 * 16 skipped lines per 1024-byte block matches the FORTH block geometry
 * exactly — if every line is unresolvable the capsule is fundamentally
 * broken.
 */
#define CAPSULE_MAX_DEFERRED CAPSULE_BLOCK_LINES

#ifdef __STARKERNEL__
/*===========================================================================
 * capsule_blk_init - Initialize block subsystem and attach the kernel ramdrive.
 *
 * Sets up the unified block address space:
 *   LBN 0..2047:                  fast RAM (ram_buf)
 *   LBN 2048..2048+KRD_MAX_BLOCKS: ramdrive (krd_buf, volatile)
 *
 * krd_buf must be pre-zeroed; KRD_MAX_BLOCKS 1 KiB slots are registered
 * as a raw device via blk_subsys_add_raw_device().
 * Disk device (if any) is attached separately after this call.
 *===========================================================================*/
#define KRD_MAX_BLOCKS 1024u

int capsule_blk_init(void *vm, uint8_t *ram_buf, size_t ram_size, uint8_t *krd_buf)
{
    int rc;
    (void)vm;   /* the block subsystem no longer takes one */
    rc = blk_subsys_init(ram_buf, ram_size);
    if (rc != 0) return rc;
    return blk_subsys_add_raw_device(krd_buf, KRD_MAX_BLOCKS);
}
#endif /* __STARKERNEL__ */

/*===========================================================================
 * Internal: header detection
 *
 * Tests whether the bytes at p begin a "Block <num>" header line.
 * On match writes the block number to *out_num and the byte position
 * immediately after the header's newline to *out_after.
 * Returns 1 on match, 0 otherwise.
 *===========================================================================*/

static int is_block_header(const uint8_t *p, const uint8_t *end,
                            uint32_t *out_num, const uint8_t **out_after)
{
    if ((size_t)(end - p) < 7u) return 0;

    if (p[0] != 'B') return 0;
    if (p[1] != 'l') return 0;
    if (p[2] != 'o') return 0;
    if (p[3] != 'c') return 0;
    if (p[4] != 'k') return 0;
    if (p[5] != ' ') return 0;

    const uint8_t *q = p + 6;
    if (q >= end || *q < '0' || *q > '9') return 0;

    uint32_t num = 0;
    while (q < end && *q >= '0' && *q <= '9') {
        num = num * 10u + (uint32_t)(*q - '0');
        q++;
    }

    while (q < end && *q != '\n') q++;
    if (q < end) q++;

    *out_num   = num;
    *out_after = q;
    return 1;
}

/*===========================================================================
 * Internal: write one block to the ramdrive (device region, 2048+)
 *===========================================================================*/

static void write_ramdrive_block(uint32_t block_num,
                                 const uint8_t *content,
                                 uint32_t content_len)
{
    uint8_t *buf = blk_get_buffer(block_num, 1);
    if (!buf) return;

    uint32_t i;
    for (i = 0; i < LOADER_BLOCK_SIZE; i++) buf[i] = 0;

    uint32_t copy_len = (content_len > LOADER_BLOCK_SIZE)
                        ? LOADER_BLOCK_SIZE : content_len;
    for (i = 0; i < copy_len; i++) buf[i] = content[i];

    blk_update(block_num);
}


/*===========================================================================
 * Internal: uint32 to decimal string
 * Returns number of characters written (not including null terminator).
 *===========================================================================*/

static size_t u32_to_dec(uint32_t val, char *buf, size_t size)
{
    char tmp[11];
    size_t len = 0;
    if (size < 2u) return 0;
    if (val == 0u) {
        tmp[len++] = '0';
    } else {
        while (val > 0u && len < 10u) {
            tmp[len++] = '0' + (char)(val % 10u);
            val /= 10u;
        }
    }
    size_t out = 0;
    while (len > 0u && out < size - 1u) buf[out++] = tmp[--len];
    buf[out] = '\0';
    return out;
}

/*===========================================================================
 * capsule_load_blocks
 *===========================================================================*/

int capsule_load_blocks(const uint8_t *payload, uint64_t length,
                        uint32_t *out_entry_block)
{
    if (!payload || length == 0) return -1;

    if (out_entry_block) *out_entry_block = 0;

    const uint8_t *p          = payload;
    const uint8_t *end        = payload + length;
    int            blocks_written    = 0;
    int            in_block          = 0;
    int            entry_set         = 0;
    uint32_t       current_block_num = 0;
    const uint8_t *block_start       = (void *)0;

    while (p < end) {
        uint32_t       block_num;
        const uint8_t *content_start;

        if (is_block_header(p, end, &block_num, &content_start)) {
            if (in_block) {
                uint32_t clen = (uint32_t)(p - block_start);
                write_ramdrive_block(current_block_num, block_start, clen);
                blocks_written++;
            }
            current_block_num = block_num;
            block_start       = content_start;
            in_block          = 1;
            if (!entry_set && out_entry_block) {
                *out_entry_block = block_num;
                entry_set = 1;
            }
            p = content_start;
        } else {
            while (p < end && *p != '\n') p++;
            if (p < end) p++;
        }
    }

    if (in_block) {
        uint32_t clen = (uint32_t)(end - block_start);
        write_ramdrive_block(current_block_num, block_start, clen);
        blocks_written++;
    }

    return blocks_written;
}

/*===========================================================================
 * exec_block_with_retry — execute one block with forward-reference tolerance
 *
 * Executes the block line by line.  Lines 0-15 (the canonical FORTH block
 * geometry) can be deferred via a uint16_t bitmask; lines >= 16 execute
 * unconditionally — an error there is non-deferrable and fails the block.
 *
 * On any error in lines 0-15 the offending line is recorded in the bitmask
 * and the entire block is retried from the top.  Deferred lines are skipped
 * in O(1) on every subsequent pass.
 *
 * Terminates when:
 *   a) block executes without error → returns 0 (deferred lines are
 *      forward references to be resolved later; logged but not fatal)
 *   b) deferred_mask == 0xFFFF (all 16 bitmask slots exhausted) → returns -1
 *   c) an error occurs on line >= 16 (beyond defer window) → returns -1
 *
 * VM compile/interpret mode is reset to MODE_INTERPRET before each retry
 * so a half-compiled colon definition cannot corrupt the next attempt.
 *===========================================================================*/

static int popcount16(uint16_t v)
{
    int c = 0;
    while (v) { c += (int)(v & 1u); v >>= 1; }
    return c;
}

static int exec_block_with_retry(VM *vm, const uint8_t *block_start,
                                  size_t block_len, uint32_t block_num)
{
    uint16_t deferred_mask = 0u;
    char     deferred_text[CAPSULE_BLOCK_LINES][CAPSULE_LINE_COLS + 1];
    int      j;

    for (j = 0; j < CAPSULE_BLOCK_LINES; j++) deferred_text[j][0] = '\0';

    for (;;) {
        /* Reset interpret/compile state before each attempt */
        vm->error     = 0;
        vm->mode      = MODE_INTERPRET;
        vm->state_var = 0;
        vm_store_cell(vm, vm->state_addr, 0);

        int            error_this_pass = 0;
        int            error_line      = 0;   /* set before use; init silences compiler */
        const uint8_t *lp              = block_start;
        const uint8_t *lend            = block_start + block_len;
        int            line_index      = 0;

        while (lp < lend) {
            const uint8_t *nl = lp;
            while (nl < lend && *nl != '\n') nl++;

            /* Only lines 0-15 can be deferred via the bitmask */
            int in_defer_range = (line_index < CAPSULE_BLOCK_LINES);

            /* O(1) skip check via bitmask — skipped only if in range AND bit set */
            if (!in_defer_range || !(deferred_mask & (1u << (unsigned)line_index))) {
                char   line[CAPSULE_EXEC_LINE_MAX];
                size_t line_len = (size_t)(nl - lp);
                size_t i;

                /* Clamp only for the execution buffer; block geometry enforced
                   at storage time, not here */
                if (line_len >= CAPSULE_EXEC_LINE_MAX)
                    line_len = CAPSULE_EXEC_LINE_MAX - 1u;
                for (i = 0; i < line_len; i++) line[i] = (char)lp[i];
                line[line_len] = '\0';

                if (line_len > 0) {
                    vm_interpret(vm, line);

                    /* ABORT: stop this block's remaining lines cleanly.
                     * Not a failure (system_word_abort() clears vm->error),
                     * so return 0 rather than -1 -- capsule_exec_payload
                     * treats any non-zero return as fatal and would stop
                     * loading the rest of the capsule payload, which would
                     * silently break word definitions in later blocks that
                     * have nothing to do with why this one aborted. */
                    if (vm->abort_requested) {
                        vm->abort_requested = 0;
                        vm->error           = 0;
                        vm->mode            = MODE_INTERPRET;
                        vm->state_var       = 0;
                        vm_store_cell(vm, vm->state_addr, 0);
                        return 0;
                    }

                    if (vm->error) {
                        vm->error     = 0;
                        vm->mode      = MODE_INTERPRET;
                        vm->state_var = 0;
                        vm_store_cell(vm, vm->state_addr, 0);
                        if (!in_defer_range) {
                            /* Beyond bitmask window — cannot defer, block fails */
#ifdef __STARKERNEL__
                            {
                                char nbuf[12], lbuf[4];
                                u32_to_dec(block_num, nbuf, sizeof(nbuf));
                                u32_to_dec((uint32_t)line_index, lbuf, sizeof(lbuf));
                                console_puts("[CAPSULE][FAIL] block ");
                                console_puts(nbuf);
                                console_puts(" line ");
                                console_puts(lbuf);
                                console_puts(": error beyond defer window\r\n");
                            }
#endif
                            return -1;
                        }
                        error_this_pass = 1;
                        error_line      = line_index;
                        /* Save truncated text for diagnostics */
                        if (deferred_text[line_index][0] == '\0') {
                            size_t tlen = line_len < CAPSULE_LINE_COLS
                                          ? line_len : CAPSULE_LINE_COLS;
                            for (i = 0; i < tlen; i++)
                                deferred_text[line_index][i] = line[i];
                            deferred_text[line_index][tlen] = '\0';
                        }
                        break;
                    }
                }
            }

            lp = (nl < lend) ? nl + 1 : lend;
            line_index++;
        }

        if (!error_this_pass) {
            /* Block passed — log any deferred lines as unresolved forward refs */
#ifdef __STARKERNEL__
            if (deferred_mask != 0u) {
                char nbuf[12], dbuf[4];
                u32_to_dec(block_num, nbuf, sizeof(nbuf));
                u32_to_dec((uint32_t)popcount16(deferred_mask), dbuf, sizeof(dbuf));
                console_puts("[CAPSULE][DEFER] block ");
                console_puts(nbuf);
                console_puts(": ");
                console_puts(dbuf);
                console_puts(" forward ref(s) unresolved at load time:\r\n");
                for (j = 0; j < CAPSULE_BLOCK_LINES; j++) {
                    if (deferred_mask & (1u << (unsigned)j)) {
                        char lbuf[4];
                        u32_to_dec((uint32_t)j, lbuf, sizeof(lbuf));
                        console_puts("  line ");
                        console_puts(lbuf);
                        console_puts(": ");
                        console_puts(deferred_text[j]);
                        console_puts("\r\n");
                    }
                }
            }
#endif
            /* Always leave interpreter in interpret mode after a block completes.
             * A deferred ';' in a nested capsule (e.g. doe.4th block 2107) can
             * strand the VM in compile mode; resetting here prevents the caller's
             * next ':' from seeing a spurious "nested colon" error. */
            vm->error     = 0;
            vm->mode      = MODE_INTERPRET;
            vm->state_var = 0;
            vm_store_cell(vm, vm->state_addr, 0);
            return 0;
        }

        /* Budget exhaustion: all 16 line slots are deferred */
        if (deferred_mask == 0xFFFFu) {
#ifdef __STARKERNEL__
            {
                char nbuf[12];
                u32_to_dec(block_num, nbuf, sizeof(nbuf));
                console_puts("[CAPSULE][PANIC] block ");
                console_puts(nbuf);
                console_puts(": exceeded 16 deferred lines — capsule cannot load\r\n");
                for (j = 0; j < CAPSULE_BLOCK_LINES; j++) {
                    if (deferred_text[j][0] != '\0') {
                        char lbuf[4];
                        u32_to_dec((uint32_t)j, lbuf, sizeof(lbuf));
                        console_puts("  line ");
                        console_puts(lbuf);
                        console_puts(": ");
                        console_puts(deferred_text[j]);
                        console_puts("\r\n");
                    }
                }
            }
#endif
            return -1;
        }

        /* Defer the errored line — set its bit */
        deferred_mask |= (1u << (unsigned)error_line);

#ifdef __STARKERNEL__
        {
            char nbuf[12], lbuf[4], dbuf[4];
            u32_to_dec(block_num, nbuf, sizeof(nbuf));
            u32_to_dec((uint32_t)error_line, lbuf, sizeof(lbuf));
            u32_to_dec((uint32_t)popcount16(deferred_mask), dbuf, sizeof(dbuf));
            console_puts("[CAPSULE][DEFER] block ");
            console_puts(nbuf);
            console_puts(" (");
            console_puts(dbuf);
            console_puts("/16) line ");
            console_puts(lbuf);
            console_puts(": ");
            console_puts(deferred_text[error_line]);
            console_puts("\r\n");
        }
#endif
        /* Retry the block with this line now in the skip bitmask */
    }
}

/*===========================================================================
 * capsule_exec_payload — unified block parse + populate + execute
 *
 * For each "Block <num>" section in the payload:
 *   1. Write content to block device; warn and truncate if > 1KB
 *   2. Execute content line-by-line via exec_block_with_retry
 *
 * Block numbers are arbitrary; order is irrelevant.
 * No LOAD is injected — FORTH code calls LOAD itself if needed.
 *===========================================================================*/

int capsule_exec_payload(void *vm_opaque, const uint8_t *payload, uint64_t length)
{
    VM *vm = (VM *)vm_opaque;
    if (!vm || !payload || length == 0) return -1;

    const uint8_t *p   = payload;
    const uint8_t *end = payload + length;

    int            in_block          = 0;
    uint32_t       current_block_num = 0;
    const uint8_t *block_start       = (void *)0;

    for (;;) {
        uint32_t       next_num      = 0;
        const uint8_t *content_start = (void *)0;
        int            at_end        = (p >= end);
        int            is_hdr        = !at_end &&
                                       is_block_header(p, end, &next_num, &content_start);

        if (in_block && (is_hdr || at_end)) {
            /* Flush completed block */
            const uint8_t *block_end = at_end ? end : p;
            size_t         block_len = (size_t)(block_end - block_start);

            if (block_len > LOADER_BLOCK_SIZE) {
#ifdef __STARKERNEL__
                char nbuf[12];
                u32_to_dec(current_block_num, nbuf, sizeof(nbuf));
                console_puts("WARN: block ");
                console_puts(nbuf);
                console_puts(" exceeds 1KB, truncating\n");
#endif
                block_len = LOADER_BLOCK_SIZE;
            }

            /* 1. Populate ramdrive (needed if FORTH later calls LOAD N) */
            write_ramdrive_block(current_block_num, block_start, (uint32_t)block_len);

            /* 2. Execute with forward-reference retry (up to 16 deferred lines) */
            log_message(LOG_INFO, "exec_block: block %u len %zu", current_block_num, block_len);
            if (exec_block_with_retry(vm, block_start, block_len,
                                      current_block_num) != 0)
                return -1;
        }

        if (at_end) break;

        if (is_hdr) {
            current_block_num = next_num;
            block_start       = content_start;
            in_block          = 1;
            p                 = content_start;
        } else {
            /* Pre-block or unrecognised line — skip */
            while (p < end && *p != '\n') p++;
            if (p < end) p++;
        }
    }

    return 0;
}

/*===========================================================================
 * capsule_clear_blocks — zero dedicated RAM slots (N-2048) for userspace
 *===========================================================================*/

void capsule_clear_blocks(const uint8_t *payload, uint64_t length)
{
    if (!payload || length == 0) return;

    const uint8_t *p   = payload;
    const uint8_t *end = payload + length;

    while (p < end) {
        uint32_t       block_num;
        const uint8_t *content_start;

        if (is_block_header(p, end, &block_num, &content_start)) {
            uint8_t *buf = blk_get_buffer(block_num, 1);
            if (buf) {
                uint32_t i;
                for (i = 0; i < LOADER_BLOCK_SIZE; i++) buf[i] = 0;
                blk_update(block_num);
            }
            p = content_start;
        } else {
            while (p < end && *p != '\n') p++;
            if (p < end) p++;
        }
    }
}

/*===========================================================================
 * capsule_exec_init — full load / execute / clear sequence
 *===========================================================================*/

CapsuleRunResult capsule_exec_init(
    void                   *vm_opaque,
    const char             *capsule_name,
    const CapsuleDirHeader *dir,
    const CapsuleDesc      *descs,
    const CapsuleNameEntry *names,
    const uint8_t          *arena)
{
    if (!vm_opaque || !capsule_name || !dir || !descs || !names || !arena)
        return CAPSULE_RUN_ERR_INVALID;

    VM *vm = (VM *)vm_opaque;

    /* Locate and validate capsule */
    const CapsuleDesc *cap = capsule_find_by_name(dir, descs, names, capsule_name);
    if (!cap) return CAPSULE_RUN_ERR_INVALID;

    CapsuleValidateResult vr = capsule_validate(cap, arena, dir->arena_size, 1);
    if (vr != CAPSULE_VALID) return CAPSULE_RUN_ERR_INVALID;

    const uint8_t *payload = capsule_get_payload(cap, arena);
    if (!payload) return CAPSULE_RUN_ERR_INVALID;

    /* Populate block device + execute. Block content is left in place
     * afterward -- capsule_exec_payload() already commits it to real
     * ramdrive storage (write_ramdrive_block()) specifically so FORTH
     * code can address it by number; a Standard BLOCK/LOAD on that same
     * number must be able to read back what EXEC just ran, the same way
     * any other block-storage write persists until something explicitly
     * overwrites or blanks it. Previously this called capsule_clear_blocks()
     * here, zeroing the slots immediately after execution -- that made
     * capsule content invisible to LOAD the instant EXEC returned, which
     * is not how FORTH-79 BLOCK/LOAD is supposed to behave, and is also
     * the reason a FORTH-79/83-locked identity (no EXEC in its own
     * dictionary, see capsules/acl-std79.4th) had no standard-compliant
     * way to ever load capsule-authored source: LOAD was on the allowlist
     * but nothing was ever left for it to find. Found and fixed 2026-09-10. */
    int rc = capsule_exec_payload(vm, payload, cap->length);
    return (rc == 0) ? CAPSULE_RUN_OK : CAPSULE_RUN_ERR_EXEC_FAIL;
}