// codegen.c -- DVX BASIC p-code emitter implementation

#include "codegen.h"
#include "symtab.h"
#include "opcodes.h"
#include "thirdparty/stb_ds_wrap.h"

#include <stdlib.h>
#include <string.h>
#include <strings.h>

// ============================================================
// basAddData
// ============================================================

bool basAddData(BasCodeGenT *cg, BasValueT val) {
    BasValueT copy = basValCopy(val);
    arrput(cg->dataPool, copy);
    cg->dataCount = (int32_t)arrlen(cg->dataPool);
    return true;
}


// ============================================================
// basAddConstant
// ============================================================

uint16_t basAddConstant(BasCodeGenT *cg, const char *text, int32_t len) {
    // Check if this string is already in the pool
    for (int32_t i = 0; i < cg->constCount; i++) {
        if (cg->constants[i]->len == len && memcmp(cg->constants[i]->data, text, len) == 0) {
            return (uint16_t)i;
        }
    }

    uint16_t idx = (uint16_t)cg->constCount;
    BasStringT *s = basStringNew(text, len);
    arrput(cg->constants, s);
    cg->constCount = (int32_t)arrlen(cg->constants);
    return idx;
}


// ============================================================
// basCodeGenBuildModule
// ============================================================

BasModuleT *basCodeGenBuildModule(BasCodeGenT *cg) {
    BasModuleT *mod = (BasModuleT *)calloc(1, sizeof(BasModuleT));

    if (!mod) {
        return NULL;
    }

    // Copy code
    mod->code = (uint8_t *)malloc(cg->codeLen);

    if (!mod->code) {
        free(mod);
        return NULL;
    }

    memcpy(mod->code, cg->code, cg->codeLen);
    mod->codeLen = cg->codeLen;

    // Copy constant pool (share string refs)
    if (cg->constCount > 0) {
        mod->constants = (BasStringT **)malloc(cg->constCount * sizeof(BasStringT *));

        if (!mod->constants) {
            free(mod->code);
            free(mod);
            return NULL;
        }

        for (int32_t i = 0; i < cg->constCount; i++) {
            mod->constants[i] = basStringRef(cg->constants[i]);
        }
    }

    mod->constCount  = cg->constCount;
    mod->globalCount = cg->globalCount;
    mod->entryPoint  = 0;

    // Copy data pool
    if (cg->dataCount > 0) {
        mod->dataPool = (BasValueT *)malloc(cg->dataCount * sizeof(BasValueT));

        if (!mod->dataPool) {
            free(mod->constants);
            free(mod->code);
            free(mod);
            return NULL;
        }

        for (int32_t i = 0; i < cg->dataCount; i++) {
            mod->dataPool[i] = basValCopy(cg->dataPool[i]);
        }
    }

    mod->dataCount = cg->dataCount;

    // Copy form variable info
    if (cg->formVarInfoCount > 0) {
        mod->formVarInfo = (BasFormVarInfoT *)malloc(cg->formVarInfoCount * sizeof(BasFormVarInfoT));

        if (mod->formVarInfo) {
            memcpy(mod->formVarInfo, cg->formVarInfo, cg->formVarInfoCount * sizeof(BasFormVarInfoT));
        }

        mod->formVarInfoCount = cg->formVarInfoCount;
    }

    return mod;
}


// ============================================================
// basCodeGenBuildModuleWithProcs
// ============================================================

BasModuleT *basCodeGenBuildModuleWithProcs(BasCodeGenT *cg, void *symtab) {
    BasModuleT *mod = basCodeGenBuildModule(cg);

    if (!mod || !symtab) {
        return mod;
    }

    BasSymTabT *tab = (BasSymTabT *)symtab;

    // Count SUB/FUNCTION entries
    int32_t procCount = 0;

    for (int32_t i = 0; i < tab->count; i++) {
        BasSymbolT *s = &tab->symbols[i];

        if ((s->kind == SYM_SUB || s->kind == SYM_FUNCTION) && s->isDefined && !s->isExtern) {
            procCount++;
        }
    }

    if (procCount == 0) {
        return mod;
    }

    mod->procs = (BasProcEntryT *)malloc(procCount * sizeof(BasProcEntryT));

    if (!mod->procs) {
        return mod;
    }

    int32_t idx = 0;

    for (int32_t i = 0; i < tab->count; i++) {
        BasSymbolT *s = &tab->symbols[i];

        if ((s->kind == SYM_SUB || s->kind == SYM_FUNCTION) && s->isDefined && !s->isExtern) {
            BasProcEntryT *p = &mod->procs[idx++];
            strncpy(p->name, s->name, BAS_MAX_PROC_NAME - 1);
            p->name[BAS_MAX_PROC_NAME - 1] = '\0';
            p->codeAddr   = s->codeAddr;
            p->paramCount = s->paramCount;
            p->returnType = s->dataType;
            p->isFunction = (s->kind == SYM_FUNCTION);
        }
    }

    mod->procCount = idx;
    return mod;
}


// ============================================================
// basCodeGenFree
// ============================================================

void basCodeGenFree(BasCodeGenT *cg) {
    for (int32_t i = 0; i < cg->constCount; i++) {
        basStringUnref(cg->constants[i]);
    }

    for (int32_t i = 0; i < cg->dataCount; i++) {
        basValRelease(&cg->dataPool[i]);
    }

    arrfree(cg->code);
    arrfree(cg->constants);
    arrfree(cg->dataPool);
    arrfree(cg->formVarInfo);
    cg->code             = NULL;
    cg->constants        = NULL;
    cg->dataPool         = NULL;
    cg->formVarInfo      = NULL;
    cg->constCount       = 0;
    cg->dataCount        = 0;
    cg->codeLen          = 0;
    cg->formVarInfoCount = 0;
}


// ============================================================
// basCodeGenInit
// ============================================================

void basCodeGenInit(BasCodeGenT *cg) {
    memset(cg, 0, sizeof(*cg));
}


// ============================================================
// basCodePos
// ============================================================

int32_t basCodePos(const BasCodeGenT *cg) {
    return cg->codeLen;
}


// ============================================================
// basEmit8
// ============================================================

void basEmit8(BasCodeGenT *cg, uint8_t b) {
    arrput(cg->code, b);
    cg->codeLen = (int32_t)arrlen(cg->code);
}


// ============================================================
// basEmit16
// ============================================================

void basEmit16(BasCodeGenT *cg, int16_t v) {
    uint8_t buf[2];
    memcpy(buf, &v, 2);
    arrput(cg->code, buf[0]);
    arrput(cg->code, buf[1]);
    cg->codeLen = (int32_t)arrlen(cg->code);
}


// ============================================================
// basEmitDouble
// ============================================================

void basEmitDouble(BasCodeGenT *cg, double v) {
    uint8_t buf[sizeof(double)];
    memcpy(buf, &v, sizeof(double));

    for (int32_t i = 0; i < (int32_t)sizeof(double); i++) {
        arrput(cg->code, buf[i]);
    }

    cg->codeLen = (int32_t)arrlen(cg->code);
}


// ============================================================
// basEmitFloat
// ============================================================

void basEmitFloat(BasCodeGenT *cg, float v) {
    uint8_t buf[sizeof(float)];
    memcpy(buf, &v, sizeof(float));

    for (int32_t i = 0; i < (int32_t)sizeof(float); i++) {
        arrput(cg->code, buf[i]);
    }

    cg->codeLen = (int32_t)arrlen(cg->code);
}


// ============================================================
// basEmitU16
// ============================================================

void basEmitU16(BasCodeGenT *cg, uint16_t v) {
    uint8_t buf[2];
    memcpy(buf, &v, 2);
    arrput(cg->code, buf[0]);
    arrput(cg->code, buf[1]);
    cg->codeLen = (int32_t)arrlen(cg->code);
}


// ============================================================
// basModuleFindProc
// ============================================================

const BasProcEntryT *basModuleFindProc(const BasModuleT *mod, const char *name) {
    if (!mod || !mod->procs || !name) {
        return NULL;
    }

    for (int32_t i = 0; i < mod->procCount; i++) {
        if (strcasecmp(mod->procs[i].name, name) == 0) {
            return &mod->procs[i];
        }
    }

    return NULL;
}


// ============================================================
// basModuleFree
// ============================================================

void basModuleFree(BasModuleT *mod) {
    if (!mod) {
        return;
    }

    free(mod->code);

    if (mod->constants) {
        for (int32_t i = 0; i < mod->constCount; i++) {
            basStringUnref(mod->constants[i]);
        }

        free(mod->constants);
    }

    if (mod->dataPool) {
        for (int32_t i = 0; i < mod->dataCount; i++) {
            basValRelease(&mod->dataPool[i]);
        }

        free(mod->dataPool);
    }

    free(mod->procs);
    free(mod->formVarInfo);
    free(mod);
}


// ============================================================
// basPatch16
// ============================================================

void basPatch16(BasCodeGenT *cg, int32_t pos, int16_t val) {
    if (pos >= 0 && pos + 2 <= cg->codeLen) {
        memcpy(&cg->code[pos], &val, 2);
    }
}