/* spl_ir.c — SIR binary serialization, deserialization, and utilities
 *
 * Binary format (all metadata fields are spl_val_t = uint64_t LE):
 *   [HEADER]     magic(8) nfuncs(8) ninsns(8) nnatives(8) nstrs(8) ndata(8)
 *   [FUNCS]      each: name_len(8) name(padded to 8) idx_of_strtab(8)
 *                       nargs(8) ninsns(8) address(8)
 *   [INSTRS]     each: opcode(2) type(2) imm(8) = 12 bytes
 *   [NATIVES]    each: name_len(8) name(padded to 8) idx_of_strtab(8)
 *   [STRTAB]     each: slen(8) str(slen bytes, padded to 8)
 */

#include "spl_ir.h"

void spl_prog_init(spl_prog_t *prog) {
    if (!prog)
        return;
    vec_init(prog->insns);
    vec_init(prog->funcs);
    vec_init(prog->natives);
    vec_init(prog->strtab);
    vec_init(prog->gdata);
    map_init(prog->symtab, MAP_HASH_STR, MAP_CMP_STR);
}

void spl_prog_drop(spl_prog_t *prog) {
    if (!prog)
        return;
    vec_free(prog->insns);
    /* free each func name */
    vec_for(prog->funcs, i) { free(vec_at(prog->funcs, i).name); }
    vec_free(prog->funcs);
    /* free each native name */
    vec_for(prog->natives, i) { free(vec_at(prog->natives, i).name); }
    vec_free(prog->natives);
    /* free each gdata entry */
    vec_for(prog->gdata, i) { free(vec_at(prog->gdata, i).data); }
    vec_free(prog->gdata);
    /* free each strtab entry */
    vec_for(prog->strtab, i) { free((void *)vec_at(prog->strtab, i)); }
    vec_free(prog->strtab);
    map_free(prog->symtab);
}

/* ---- LE read/write helpers ---- */

static inline spl_val_t rd64(const unsigned char **p) {
    spl_val_t v = (spl_val_t)(*p)[0] | ((spl_val_t)(*p)[1] << 8) | ((spl_val_t)(*p)[2] << 16) |
                  ((spl_val_t)(*p)[3] << 24) | ((spl_val_t)(*p)[4] << 32) |
                  ((spl_val_t)(*p)[5] << 40) | ((spl_val_t)(*p)[6] << 48) |
                  ((spl_val_t)(*p)[7] << 56);
    *p += 8;
    return v;
}

static inline void wr64(unsigned char **p, spl_val_t v) {
    *(*p)++ = (unsigned char)(v);
    *(*p)++ = (unsigned char)(v >> 8);
    *(*p)++ = (unsigned char)(v >> 16);
    *(*p)++ = (unsigned char)(v >> 24);
    *(*p)++ = (unsigned char)(v >> 32);
    *(*p)++ = (unsigned char)(v >> 40);
    *(*p)++ = (unsigned char)(v >> 48);
    *(*p)++ = (unsigned char)(v >> 56);
}

/* Round n up to next multiple of 8 */
#define ALIGN8(n) (((n) + 7) & ~7)

int spl_prog_load_from_file(const char *fname, spl_prog_t *prog) {
    FILE *f;
    unsigned char *data;
    long len;
    const unsigned char *p;
    spl_val_t nfuncs, ninsns, nnatives, nstrs, ndata;

    if (!fname || !prog)
        return -1;

    f = fopen(fname, "rb");
    if (!f)
        return -1;
    fseek(f, 0, SEEK_END);
    len = ftell(f);
    fseek(f, 0, SEEK_SET);
    if (len < 48) {
        fclose(f);
        return -1;
    }
    data = (unsigned char *)malloc((size_t)len);
    if (!data) {
        fclose(f);
        return -1;
    }
    if (fread(data, 1, (size_t)len, f) != (size_t)len) {
        free(data);
        fclose(f);
        return -1;
    }
    fclose(f);

    p = data;

    /* magic */
    if (p[0] != 'S' || p[1] != 'P' || p[2] != 'L' || p[3] != 'B' || p[4] != 'I' || p[5] != 'N' ||
        p[6] != '\0' || p[7] != '\0') {
        free(data);
        return -1;
    }
    p += 8;

    spl_prog_init(prog);

    /* header counts */
    nfuncs = rd64(&p);
    ninsns = rd64(&p);
    nnatives = rd64(&p);
    nstrs = rd64(&p);
    ndata = rd64(&p);

    /* ---- function table ---- */
    for (spl_val_t i = 0; i < nfuncs; i++) {
        spl_func_t func = {0};
        spl_val_t nlen = rd64(&p);
        usize pad = ALIGN8((usize)nlen) - (usize)nlen;

        func.name = (char *)malloc((usize)nlen);
        if (!func.name) {
            free(data);
            return -1;
        }
        memcpy(func.name, p, (usize)nlen);
        p += (usize)nlen + pad;

        func.idx_of_strtab = rd64(&p);
        func.nargs = rd64(&p);
        func.ninsns = rd64(&p);
        func.address = rd64(&p);
        vec_push(prog->funcs, func);
    }

    /* ---- instructions ---- */
    for (spl_val_t i = 0; i < ninsns; i++) {
        if ((size_t)(p - data) + 12 > (size_t)len) {
            free(data);
            return -1;
        }
        spl_ins_t ins;
        ins.opcode = (uint16_t)p[0] | ((uint16_t)p[1] << 8);
        ins.type = (uint16_t)p[2] | ((uint16_t)p[3] << 8);
        p += 4;
        ins.imm = rd64(&p);
        vec_push(prog->insns, ins);
    }

    /* ---- native table ---- */
    for (spl_val_t i = 0; i < nnatives; i++) {
        spl_native_t nat = {0};
        spl_val_t nlen = rd64(&p);
        usize pad = ALIGN8((usize)nlen) - (usize)nlen;

        nat.name = (char *)malloc((usize)nlen);
        if (!nat.name) {
            free(data);
            return -1;
        }
        memcpy(nat.name, p, (usize)nlen);
        p += (usize)nlen + pad;

        nat.idx_of_strtab = rd64(&p);
        nat.impl_fn = NULL; /* function pointer can't be serialised */
        vec_push(prog->natives, nat);
    }

    /* ---- string table ---- */
    for (spl_val_t i = 0; i < nstrs; i++) {
        spl_val_t slen = rd64(&p);
        usize pad = ALIGN8((usize)slen) - (usize)slen;
        char *s = (char *)malloc((usize)slen + 1);
        if (!s) {
            free(data);
            return -1;
        }
        memcpy(s, p, (usize)slen);
        s[(usize)slen] = '\0';
        p += (usize)slen + pad;
        vec_push(prog->strtab, s);
    }

    /* ---- global data ---- */
    for (spl_val_t i = 0; i < ndata; i++) {
        spl_gdata_t entry;
        entry.size = rd64(&p);
        usize pad = ALIGN8(entry.size) - entry.size;
        entry.data = (unsigned char *)malloc(entry.size);
        if (!entry.data) {
            free(data);
            return -1;
        }
        memcpy(entry.data, p, entry.size);
        p += entry.size + pad;
        vec_push(prog->gdata, entry);
    }

    free(data);
    return 0;
}

int spl_prog_store_to_file(const char *fname, spl_prog_t *prog) {
    unsigned char *buf, *p;
    spl_val_t i;
    size_t total;
    spl_val_t nfuncs, ninsns, nnatives, nstrs, ndata;
    usize nlen, pad;

    if (!fname || !prog)
        return -1;

    nfuncs = vec_size(prog->funcs);
    ninsns = vec_size(prog->insns);
    nnatives = vec_size(prog->natives);
    nstrs = vec_size(prog->strtab);
    ndata = vec_size(prog->gdata);

    /* Calculate total size */
    total = 8                    /* magic */
            + 8 + 8 + 8 + 8 + 8; /* 5 counts (nfuncs+ninsns+nnatives+nstrs+ndata) */

    /* funcs */
    for (i = 0; i < nfuncs; i++) {
        nlen = strlen(vec_at(prog->funcs, i).name) + 1; /* include null */
        total +=
            8 + ALIGN8(nlen) + 8 + 8 + 8 + 8; /* nlen + name(pad) + idx + nargs + ninsns + addr */
    }

    /* insns */
    total += ninsns * 12; /* opcode(2) + type(2) + imm(8) */

    /* natives */
    for (i = 0; i < nnatives; i++) {
        nlen = strlen(vec_at(prog->natives, i).name) + 1;
        total += 8 + ALIGN8(nlen) + 8; /* nlen + name(pad) + idx_of_strtab */
    }

    /* strtab */
    for (i = 0; i < nstrs; i++) {
        nlen = strlen(vec_at(prog->strtab, i));
        total += 8 + ALIGN8(nlen); /* slen + str(pad) */
    }

    /* gdata */
    for (i = 0; i < ndata; i++) {
        usize dsize = vec_at(prog->gdata, i).size;
        total += 8 + ALIGN8(dsize); /* dsize + data(pad) */
    }

    buf = (unsigned char *)malloc(total);
    if (!buf)
        return -1;
    p = buf;

    /* magic */
    memcpy(p, "SPLBIN\0\0", 8);
    p += 8;

    /* counts */
    wr64(&p, nfuncs);
    wr64(&p, ninsns);
    wr64(&p, nnatives);
    wr64(&p, nstrs);
    wr64(&p, ndata);

    /* ---- function table ---- */
    for (i = 0; i < nfuncs; i++) {
        const char *name = vec_at(prog->funcs, i).name;
        nlen = strlen(name) + 1;
        pad = ALIGN8(nlen) - nlen;
        wr64(&p, nlen);
        memcpy(p, name, nlen);
        p += nlen;
        memset(p, 0, pad);
        p += pad;
        wr64(&p, vec_at(prog->funcs, i).idx_of_strtab);
        wr64(&p, vec_at(prog->funcs, i).nargs);
        wr64(&p, vec_at(prog->funcs, i).ninsns);
        wr64(&p, vec_at(prog->funcs, i).address);
    }

    /* ---- instructions ---- */
    for (i = 0; i < ninsns; i++) {
        spl_ins_t *ins = &vec_at(prog->insns, i);
        *p++ = (unsigned char)(ins->opcode);
        *p++ = (unsigned char)(ins->opcode >> 8);
        *p++ = (unsigned char)(ins->type);
        *p++ = (unsigned char)(ins->type >> 8);
        wr64(&p, ins->imm);
    }

    /* ---- native table ---- */
    for (i = 0; i < nnatives; i++) {
        const char *name = vec_at(prog->natives, i).name;
        nlen = strlen(name) + 1;
        pad = ALIGN8(nlen) - nlen;
        wr64(&p, nlen);
        memcpy(p, name, nlen);
        p += nlen;
        memset(p, 0, pad);
        p += pad;
        wr64(&p, vec_at(prog->natives, i).idx_of_strtab);
    }

    /* ---- string table ---- */
    for (i = 0; i < nstrs; i++) {
        const char *s = vec_at(prog->strtab, i);
        nlen = strlen(s);
        pad = ALIGN8(nlen) - nlen;
        wr64(&p, nlen);
        memcpy(p, s, nlen);
        p += nlen;
        memset(p, 0, pad);
        p += pad;
    }

    /* ---- global data ---- */
    for (i = 0; i < ndata; i++) {
        spl_gdata_t *entry = &vec_at(prog->gdata, i);
        pad = ALIGN8(entry->size) - entry->size;
        wr64(&p, entry->size);
        memcpy(p, entry->data, entry->size);
        p += entry->size;
        memset(p, 0, pad);
        p += pad;
    }

    /* Write file */
    FILE *f = fopen(fname, "wb");
    if (!f) {
        free(buf);
        return -1;
    }
    fwrite(buf, 1, total, f);
    fclose(f);
    free(buf);
    return 0;
}

int spl_prog_add_func(spl_prog_t *prog, spl_func_t *func) {
    if (!prog || !func)
        return 0;
    vec_push(prog->funcs, *func);
    return vec_size(prog->funcs);
}

int spl_prog_add_native(spl_prog_t *prog, spl_native_t *native) {
    if (!prog || !native)
        return 0;
    vec_push(prog->natives, *native);
    return vec_size(prog->natives);
}

int spl_prog_add_data(spl_prog_t *prog, void *ptr, usize size) {
    spl_gdata_t entry;
    if (!prog)
        return 0;
    entry.data = (unsigned char *)malloc(size);
    if (!entry.data)
        return 0;
    memcpy(entry.data, ptr, size);
    entry.size = size;
    vec_push(prog->gdata, entry);
    return vec_size(prog->gdata);
}

spl_val_t spl_prog_emit(spl_prog_t *prog, uint16_t opcode, uint16_t type, spl_val_t imm) {
    spl_ins_t ins = {opcode, type, imm};
    spl_val_t addr = vec_size(prog->insns);
    vec_push(prog->insns, ins);
    return addr;
}

int spl_prog_add_func_simple(spl_prog_t *prog, const char *name, spl_val_t nargs) {
    spl_func_t func = {0};
    if (!prog || !name)
        return -1;
    func.name = strdup(name);
    func.idx_of_strtab = 0;
    func.nargs = nargs;
    func.ninsns = 0;
    func.address = vec_size(prog->insns);
    vec_push(prog->funcs, func);
    return (int)vec_size(prog->funcs) - 1;
}

int spl_prog_update_func(spl_prog_t *prog, int func_idx, const char *name, spl_val_t nargs) {
    if (!prog || func_idx < 0 || func_idx >= (int)vec_size(prog->funcs))
        return spl_prog_add_func_simple(prog, name, nargs);
    spl_func_t *func = &vec_at(prog->funcs, func_idx);
    func->address = vec_size(prog->insns);
    func->nargs = nargs;
    return func_idx;
}

void spl_prog_end_func(spl_prog_t *prog, int func_idx) {
    if (!prog || func_idx < 0 || func_idx >= (int)vec_size(prog->funcs))
        return;
    spl_func_t *func = &vec_at(prog->funcs, func_idx);
    func->ninsns = vec_size(prog->insns) - func->address;
}

int spl_prog_add_str(spl_prog_t *prog, const char *str) {
    if (!prog || !str)
        return -1;
    vec_push(prog->strtab, strdup(str));
    return (int)vec_size(prog->strtab) - 1;
}

spl_func_t *spl_prog_get_func(spl_prog_t *prog, const char *name) {
    if (!prog || !name)
        return NULL;
    vec_for(prog->funcs, i) {
        const char *match_name = vec_at(prog->funcs, i).name;
        if (match_name && strcmp(match_name, name) == 0) {
            return &vec_at(prog->funcs, i);
        }
    }
    return NULL;
}

spl_native_t *spl_prog_get_native(spl_prog_t *prog, const char *name) {
    if (!prog || !name)
        return NULL;
    vec_for(prog->natives, i) {
        const char *match_name = vec_at(prog->natives, i).name;
        if (match_name && strcmp(match_name, name) == 0) {
            return &vec_at(prog->natives, i);
        }
    }
    return NULL;
}

const char *opcode_name[] = {
#define X(opcode, name, argc, pop, push, desc) [opcode] = name,
    SPL_OPCODES(X)
#undef X
};
const char *spl_opcode_name(spl_opcode_t opcode) { return opcode_name[opcode]; }
const char *spl_type_name(spl_type_t type) {
    switch (type) {
    case SPL_VOID:
        return "void";
    case SPL_I8:
        return "i8";
    case SPL_U8:
        return "u8";
    case SPL_I16:
        return "i16";
    case SPL_U16:
        return "u16";
    case SPL_I32:
        return "i32";
    case SPL_U32:
        return "u32";
    case SPL_I64:
        return "i64";
    case SPL_U64:
        return "u64";
    case SPL_F32:
        return "f32";
    case SPL_F64:
        return "f64";
    case SPL_USIZE:
        return "usize";
    case SPL_ISIZE:
        return "isize";
    case SPL_PTR:
        return "ptr";
    default:
        return "???";
    }
}

void spl_ins_dump(spl_ins_t *ins, spl_val_t addr) {
    printf("%4zu: %s", addr, spl_opcode_name((spl_opcode_t)ins->opcode));
    if (ins->type != SPL_VOID)
        printf(" %s", spl_type_name((spl_type_t)ins->type));
    printf(" %zd:%zx", ins->imm, ins->imm);
    printf("\n");
}