// The MIT License (MIT)
//
// Copyright (C) 2026 Scott Duensing
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE 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.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.

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

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

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


// Function prototypes (alphabetical)
uint16_t basAddConstant(BasCodeGenT *cg, const char *text, int32_t len);
void basAddData(BasCodeGenT *cg, BasValueT val);
void basCodeGenAddDebugVar(BasCodeGenT *cg, const char *name, uint8_t scope, uint8_t dataType, int32_t index, int32_t procIndex, const char *formName);
BasModuleT *basCodeGenBuildModule(BasCodeGenT *cg);
BasModuleT *basCodeGenBuildModuleWithProcs(BasCodeGenT *cg, void *symtab);
void basCodeGenFree(BasCodeGenT *cg);
void basCodeGenInit(BasCodeGenT *cg);
int32_t basCodePos(const BasCodeGenT *cg);
void basEmit16(BasCodeGenT *cg, int16_t v);
void basEmit8(BasCodeGenT *cg, uint8_t b);
static void basEmitBytes(BasCodeGenT *cg, const void *data, int32_t len);
void basEmitDouble(BasCodeGenT *cg, double v);
void basEmitFloat(BasCodeGenT *cg, float v);
void basEmitU16(BasCodeGenT *cg, uint16_t v);
const BasProcEntryT *basModuleFindProc(const BasModuleT *mod, const char *name);
void basPatch16(BasCodeGenT *cg, int32_t pos, int16_t val);

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;
        }
    }

    if (cg->constCount >= BAS_MAX_CONSTANTS) {
        cg->overflow = true;
        return 0;
    }

    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;
}


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


void basCodeGenAddDebugVar(BasCodeGenT *cg, const char *name, uint8_t scope, uint8_t dataType, int32_t index, int32_t procIndex, const char *formName) {
    BasDebugVarT dv;
    memset(&dv, 0, sizeof(dv));
    snprintf(dv.name, BAS_MAX_PROC_NAME, "%s", name);
    dv.scope     = scope;
    dv.dataType  = dataType;
    dv.index     = index;
    dv.procIndex = procIndex;

    if (formName && formName[0]) {
        snprintf(dv.formName, BAS_MAX_PROC_NAME, "%s", formName);
    }
    arrput(cg->debugVars, dv);
    cg->debugVarCount = (int32_t)arrlen(cg->debugVars);
}


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) {
            // Release the constant pool refs taken above before freeing.
            for (int32_t i = 0; i < cg->constCount; i++) {
                basStringUnref(mod->constants[i]);
            }

            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));
            // Keep count in sync with the pointer: on OOM the array stays
            // NULL, so leaving count nonzero would make consumers deref NULL.
            mod->formVarInfoCount = cg->formVarInfoCount;
        }
    }

    // Copy debug variable info
    if (cg->debugVarCount > 0) {
        mod->debugVars = (BasDebugVarT *)malloc(cg->debugVarCount * sizeof(BasDebugVarT));

        if (mod->debugVars) {
            memcpy(mod->debugVars, cg->debugVars, cg->debugVarCount * sizeof(BasDebugVarT));
            mod->debugVarCount = cg->debugVarCount;
        }
    }

    // Copy UDT type definitions for debugger
    if (cg->debugUdtDefCount > 0) {
        mod->debugUdtDefs = (BasDebugUdtDefT *)malloc(cg->debugUdtDefCount * sizeof(BasDebugUdtDefT));

        if (mod->debugUdtDefs) {
            for (int32_t i = 0; i < cg->debugUdtDefCount; i++) {
                mod->debugUdtDefs[i] = cg->debugUdtDefs[i];
                mod->debugUdtDefs[i].fields = NULL;

                if (cg->debugUdtDefs[i].fieldCount > 0 && cg->debugUdtDefs[i].fields) {
                    mod->debugUdtDefs[i].fields = (BasDebugFieldT *)malloc(
                        cg->debugUdtDefs[i].fieldCount * sizeof(BasDebugFieldT));

                    if (mod->debugUdtDefs[i].fields) {
                        memcpy(mod->debugUdtDefs[i].fields, cg->debugUdtDefs[i].fields,
                               cg->debugUdtDefs[i].fieldCount * sizeof(BasDebugFieldT));
                    }
                }
            }

            mod->debugUdtDefCount = cg->debugUdtDefCount;
        }
    }

    return mod;
}


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;

    // Count globals that need runtime type init (currently: STRING).
    // This list survives stripping, unlike debugVars.
    int32_t globalInitCount = 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 (s->scope == SCOPE_GLOBAL && s->kind == SYM_VARIABLE && s->dataType == BAS_TYPE_STRING && !s->isArray) {
            globalInitCount++;
        }
    }

    if (globalInitCount > 0) {
        mod->globalInits = (BasGlobalInitT *)malloc(globalInitCount * sizeof(BasGlobalInitT));

        if (mod->globalInits) {
            int32_t gi = 0;

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

                if (s->scope == SCOPE_GLOBAL && s->kind == SYM_VARIABLE && s->dataType == BAS_TYPE_STRING && !s->isArray) {
                    mod->globalInits[gi].index    = s->index;
                    mod->globalInits[gi].dataType = s->dataType;
                    gi++;
                }
            }

            mod->globalInitCount = gi;
        }
    }

    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';
            strncpy(p->formName, s->formName, BAS_MAX_PROC_NAME - 1);
            p->formName[BAS_MAX_PROC_NAME - 1] = '\0';
            p->codeAddr   = s->codeAddr;
            p->paramCount = s->paramCount;
            p->localCount = s->localCount;
            p->returnType = s->dataType;
            p->isFunction = (s->kind == SYM_FUNCTION);
        }
    }

    mod->procCount = idx;
    return mod;
}


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);
    arrfree(cg->debugVars);

    for (int32_t i = 0; i < cg->debugUdtDefCount; i++) {
        free(cg->debugUdtDefs[i].fields);
    }

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


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


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


void basEmit16(BasCodeGenT *cg, int16_t v) {
    basEmitBytes(cg, &v, (int32_t)sizeof(v));
}


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


static void basEmitBytes(BasCodeGenT *cg, const void *data, int32_t len) {
    const uint8_t *p = (const uint8_t *)data;

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

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


void basEmitDouble(BasCodeGenT *cg, double v) {
    basEmitBytes(cg, &v, (int32_t)sizeof(v));
}


void basEmitFloat(BasCodeGenT *cg, float v) {
    basEmitBytes(cg, &v, (int32_t)sizeof(v));
}


void basEmitU16(BasCodeGenT *cg, uint16_t v) {
    basEmitBytes(cg, &v, (int32_t)sizeof(v));
}


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;
}


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