// test_compact.c -- verify that strip + compact produces identical output
//
// Compiles each test program, runs it once without compaction and once
// with strip + compact applied, and compares captured PRINT output.
// Exits nonzero if any mismatch is found.

#include "compiler/parser.h"
#include "compiler/strip.h"
#include "compiler/compact.h"
#include "runtime/vm.h"
#include "runtime/values.h"

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>


#define MAX_OUT 65536


typedef struct {
    char   *buf;
    int32_t len;
    int32_t cap;
} OutBufT;


static void captureCallback(void *ctx, const char *text, bool newline) {
    OutBufT *ob   = (OutBufT *)ctx;
    int32_t  tlen = (int32_t)strlen(text);

    if (ob->len + tlen + 2 >= ob->cap) {
        return;
    }

    memcpy(ob->buf + ob->len, text, tlen);
    ob->len += tlen;

    if (newline) {
        ob->buf[ob->len++] = '\n';
    }
}


static int32_t runAndCapture(const char *source, bool compact, char *outBuf, int32_t outCap) {
    BasParserT parser;
    basParserInit(&parser, source, (int32_t)strlen(source));

    if (!basParse(&parser)) {
        fprintf(stderr, "compile error: %s\n", parser.error);
        basParserFree(&parser);
        return -1;
    }

    BasModuleT *mod = basParserBuildModule(&parser);
    basParserFree(&parser);

    if (!mod) {
        return -1;
    }

    if (compact) {
        basStripModule(mod);
        basCompactBytecode(mod);
    }

    BasVmT *vm = basVmCreate();
    basVmLoadModule(vm, mod);

    OutBufT ob;
    ob.buf = outBuf;
    ob.len = 0;
    ob.cap = outCap;

    basVmSetPrintCallback(vm, captureCallback, &ob);

    vm->callStack[0].localCount = mod->globalCount > 64 ? 64 : mod->globalCount;
    vm->callDepth = 1;

    BasVmResultE result = basVmRun(vm);

    if (result != BAS_VM_HALTED && result != BAS_VM_OK) {
        fprintf(stderr, "vm error %d: %s\n", result, basVmGetError(vm));
    }

    basVmDestroy(vm);
    basModuleFree(mod);

    outBuf[ob.len] = '\0';
    return ob.len;
}


static int32_t sTotal  = 0;
static int32_t sFailed = 0;


static void testCompact(const char *name, const char *source) {
    sTotal++;

    static char orig[MAX_OUT];
    static char comp[MAX_OUT];

    int32_t origLen = runAndCapture(source, false, orig, MAX_OUT);
    int32_t compLen = runAndCapture(source, true, comp, MAX_OUT);

    if (origLen < 0 || compLen < 0) {
        printf("FAIL: %s (runtime error)\n", name);
        sFailed++;
        return;
    }

    if (origLen != compLen || memcmp(orig, comp, origLen) != 0) {
        printf("FAIL: %s\n", name);
        printf("  uncompacted (%d bytes):\n%s\n", (int)origLen, orig);
        printf("  compacted   (%d bytes):\n%s\n", (int)compLen, comp);
        sFailed++;
        return;
    }

    printf("PASS: %s\n", name);
}


int main(void) {
    printf("DVX BASIC Bytecode Compaction Tests\n");
    printf("====================================\n\n");

    basStringSystemInit();

    // ---- Basic control flow ----

    testCompact("FOR loop",
        "DIM i AS INTEGER\n"
        "FOR i = 1 TO 5\n"
        "    PRINT i;\n"
        "NEXT i\n"
        "PRINT\n"
    );

    testCompact("Nested FOR",
        "DIM i AS INTEGER\n"
        "DIM j AS INTEGER\n"
        "FOR i = 1 TO 3\n"
        "    FOR j = 1 TO 2\n"
        "        PRINT i * 10 + j;\n"
        "    NEXT j\n"
        "NEXT i\n"
        "PRINT\n"
    );

    testCompact("FOR with negative STEP",
        "DIM i AS INTEGER\n"
        "FOR i = 5 TO 1 STEP -1\n"
        "    PRINT i;\n"
        "NEXT i\n"
        "PRINT\n"
    );

    testCompact("EXIT FOR",
        "DIM i AS INTEGER\n"
        "FOR i = 1 TO 100\n"
        "    IF i = 4 THEN EXIT FOR\n"
        "    PRINT i;\n"
        "NEXT i\n"
        "PRINT\n"
    );

    testCompact("DO WHILE",
        "DIM n AS INTEGER\n"
        "n = 0\n"
        "DO WHILE n < 3\n"
        "    PRINT n;\n"
        "    n = n + 1\n"
        "LOOP\n"
        "PRINT\n"
    );

    testCompact("IF THEN ELSE",
        "DIM x AS INTEGER\n"
        "x = 5\n"
        "IF x > 3 THEN\n"
        "    PRINT \"big\"\n"
        "ELSE\n"
        "    PRINT \"small\"\n"
        "END IF\n"
    );

    testCompact("SELECT CASE",
        "DIM g AS STRING\n"
        "g = \"B\"\n"
        "SELECT CASE g\n"
        "    CASE \"A\"\n"
        "        PRINT \"A\"\n"
        "    CASE \"B\", \"C\"\n"
        "        PRINT \"BC\"\n"
        "    CASE ELSE\n"
        "        PRINT \"?\"\n"
        "END SELECT\n"
    );

    // ---- GOSUB: exercises the absolute address in OP_PUSH_INT32 pattern ----

    testCompact("GOSUB",
        "DIM n AS INTEGER\n"
        "n = 10\n"
        "GOSUB doubler\n"
        "PRINT n\n"
        "GOSUB doubler\n"
        "PRINT n\n"
        "END\n"
        "doubler:\n"
        "n = n * 2\n"
        "RETURN\n"
    );

    testCompact("Multiple GOSUBs",
        "DIM x AS INTEGER\n"
        "x = 1\n"
        "GOSUB a\n"
        "GOSUB b\n"
        "GOSUB c\n"
        "PRINT x\n"
        "END\n"
        "a:\n"
        "x = x + 10\n"
        "RETURN\n"
        "b:\n"
        "x = x + 100\n"
        "RETURN\n"
        "c:\n"
        "x = x + 1000\n"
        "RETURN\n"
    );

    // ---- SUB / FUNCTION: exercises absolute CALL addresses + proc table ----

    testCompact("SUB with CALL",
        "CALL Greet\n"
        "CALL Greet\n"
        "CALL Greet\n"
        "SUB Greet\n"
        "    PRINT \"Hello\"\n"
        "END SUB\n"
    );

    testCompact("FUNCTION return value",
        "DECLARE FUNCTION Square(x AS INTEGER) AS INTEGER\n"
        "DIM r AS INTEGER\n"
        "r = Square(5)\n"
        "PRINT r\n"
        "r = Square(7)\n"
        "PRINT r\n"
        "FUNCTION Square(x AS INTEGER) AS INTEGER\n"
        "    Square = x * x\n"
        "END FUNCTION\n"
    );

    testCompact("Recursive FUNCTION",
        "DECLARE FUNCTION Fact(n AS INTEGER) AS INTEGER\n"
        "PRINT Fact(5)\n"
        "FUNCTION Fact(n AS INTEGER) AS INTEGER\n"
        "    IF n <= 1 THEN\n"
        "        Fact = 1\n"
        "    ELSE\n"
        "        Fact = n * Fact(n - 1)\n"
        "    END IF\n"
        "END FUNCTION\n"
    );

    // ---- ON ERROR: exercises the relative int16 offset path with non-zero ----

    testCompact("ON ERROR GOTO",
        "ON ERROR GOTO handler\n"
        "PRINT 10 / 0\n"
        "PRINT \"unreached\"\n"
        "END\n"
        "handler:\n"
        "PRINT \"caught\"\n"
    );

    // ---- Mixed: many OP_LINE between jumps ----

    testCompact("Long function with many lines",
        "DIM total AS INTEGER\n"
        "total = 0\n"
        "DIM i AS INTEGER\n"
        "FOR i = 1 TO 10\n"
        "    IF i MOD 2 = 0 THEN\n"
        "        total = total + i\n"
        "    ELSE\n"
        "        total = total + i * 2\n"
        "    END IF\n"
        "NEXT i\n"
        "PRINT total\n"
    );

    testCompact("GOSUB inside FOR",
        "DIM i AS INTEGER\n"
        "DIM sum AS INTEGER\n"
        "sum = 0\n"
        "FOR i = 1 TO 5\n"
        "    GOSUB addit\n"
        "NEXT i\n"
        "PRINT sum\n"
        "END\n"
        "addit:\n"
        "sum = sum + i\n"
        "RETURN\n"
    );

    testCompact("FOR inside SUB",
        "CALL Loop5\n"
        "SUB Loop5\n"
        "    DIM k AS INTEGER\n"
        "    FOR k = 0 TO 4\n"
        "        PRINT k;\n"
        "    NEXT k\n"
        "    PRINT\n"
        "END SUB\n"
    );

    printf("\n%d/%d tests passed\n", (int)(sTotal - sFailed), (int)sTotal);
    return sFailed > 0 ? 1 : 0;
}