/*
 * nwc, a NanoWasm compiler
 * Copyright (C) 2025 Xavier Del Campo Romero
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "instr.h"
#include "types.h"
#include <errno.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

struct imm
{
    enum {IMM_VARUINT32, IMM_TABLE} type;

    union
    {
        varuint32 varuint32;

        struct imm_table
        {
            varuint32 target_count, *targets;
        } table;
    } u;
};

static void free_imm(struct imm *const imm)
{
    switch (imm->type)
    {
        case IMM_VARUINT32:
            break;

        case IMM_TABLE:
        {
            struct imm_table *const t = &imm->u.table;

            free(t->targets);
        }
            break;
    }
}

static int imm_none(FILE *const f, struct imm *const imm)
{
    return 0;
}

static int instr_sig(FILE *const f, struct imm *const imm)
{
    return read_varint7(f, &(varint7){0});
}

static int imm_readvaruint1(FILE *const f, struct imm *const imm)
{
    return read_varuint1(f, &(varuint1){0});
}

static int imm_varuint32(FILE *const f, struct imm *const imm)
{
    imm->type = IMM_VARUINT32;
    return read_varuint32(f, &imm->u.varuint32);
}

static int imm_varuint64(FILE *const f, struct imm *const imm)
{
    return read_varuint64(f, &(varuint64){0});
}

static int imm_varint32(FILE *const f, struct imm *const imm)
{
    return read_varint32(f, &(varint32){0});
}

static int imm_varint64(FILE *const f, struct imm *const imm)
{
    return read_varint64(f, &(varint64){0});
}

static int imm_memory(FILE *const f, struct imm *const imm)
{
    return read_varuint32(f, &(varuint32){0})
        || read_varuint32(f, &(varuint32){0});
}

static int imm_uint32(FILE *const f, struct imm *const imm)
{
    uint32_t v;

    return !fread(&v, sizeof v, 1, f);
}

static int imm_uint64(FILE *const f, struct imm *const imm)
{
    uint64_t v;

    return !fread(&v, sizeof v, 1, f);
}

static int imm_table(FILE *const f, struct imm *const imm)
{
    struct imm table = {.type = IMM_TABLE};
    struct imm_table *const t = &table.u.table;

    if (read_varuint32(f, &t->target_count))
    {
        fprintf(stderr, "%s: read_varuint32 failed\n", __func__);
        goto failure;
    }
    else if (!(t->targets = malloc((t->target_count + 1) * sizeof *t->targets)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        goto failure;
    }

    for (varuint32 i = 0; i < (t->target_count + 1); i++)
        if (read_varuint32(f, &t->targets[i]))
        {
            fprintf(stderr, "%s: read_varuint32 %lu failed\n", __func__,
                (unsigned long)i);
            goto failure;
        }

    *imm = table;
    return 0;

failure:
    free_imm(&table);
    return -1;
}

static int imm_call_indirect(FILE *const f, struct imm *const imm)
{
    varuint32 reserved;

    if (read_varuint32(f, &(varuint32){0}))
    {
        fprintf(stderr, "%s: read_varuint32 failed\n", __func__);
        return -1;
    }
    else if (read_varuint32(f, &reserved))
    {
        fprintf(stderr, "%s: read_varuint1 failed\n", __func__);
        return -1;
    }
    else if (reserved)
    {
        fprintf(stderr, "%s: expected zero for reserved field\n", __func__);
        return -1;
    }

    return 0;
}

static int imm_extra(FILE *const f, struct imm *const imm)
{
    uint8_t b;

    return !fread(&b, sizeof b, 1, f);
}

#define OPS \
    X(OP_UNREACHABLE, imm_none) \
    X(OP_NOP, imm_none) \
    X(OP_BLOCK, instr_sig) \
    X(OP_LOOP, instr_sig) \
    X(OP_IF, instr_sig) \
    X(OP_ELSE, imm_none) \
    X(OP_END, imm_none) \
    X(OP_BR, imm_varuint32) \
    X(OP_BR_IF, imm_varuint32) \
    X(OP_BR_TABLE, imm_table) \
    X(OP_RETURN, imm_none) \
    X(OP_CALL, imm_varuint32) \
    X(OP_CALL_INDIRECT, imm_call_indirect) \
    X(OP_DROP, imm_none) \
    X(OP_SELECT, imm_none) \
    X(OP_GET_LOCAL, imm_varuint32) \
    X(OP_SET_LOCAL, imm_varuint32) \
    X(OP_TEE_LOCAL, imm_varuint32) \
    X(OP_GET_GLOBAL, imm_varuint32) \
    X(OP_SET_GLOBAL, imm_varuint32) \
    X(OP_I32_LOAD, imm_memory) \
    X(OP_I64_LOAD, imm_memory) \
    X(OP_F32_LOAD, imm_memory) \
    X(OP_F64_LOAD, imm_memory) \
    X(OP_I32_LOAD8_S, imm_memory) \
    X(OP_I32_LOAD8_U, imm_memory) \
    X(OP_I32_LOAD16_S, imm_memory) \
    X(OP_I32_LOAD16_U, imm_memory) \
    X(OP_I64_LOAD8_S, imm_memory) \
    X(OP_I64_LOAD8_U, imm_memory) \
    X(OP_I64_LOAD16_S, imm_memory) \
    X(OP_I64_LOAD16_U, imm_memory) \
    X(OP_I64_LOAD32_S, imm_memory) \
    X(OP_I64_LOAD32_U, imm_memory) \
    X(OP_I32_STORE, imm_memory) \
    X(OP_I64_STORE, imm_memory) \
    X(OP_F32_STORE, imm_memory) \
    X(OP_F64_STORE, imm_memory) \
    X(OP_I32_STORE8, imm_memory) \
    X(OP_I32_STORE16, imm_memory) \
    X(OP_I64_STORE8, imm_memory) \
    X(OP_I64_STORE16, imm_memory) \
    X(OP_I64_STORE132, imm_memory) \
    X(OP_CURRENT_MEMORY, imm_readvaruint1) \
    X(OP_GROW_MEMORY, imm_readvaruint1) \
    X(OP_I32_CONST, imm_varint32) \
    X(OP_I64_CONST, imm_varint64) \
    X(OP_F32_CONST, imm_uint32) \
    X(OP_F64_CONST, imm_uint64) \
    X(OP_I32_EQZ, imm_none) \
    X(OP_I32_EQ, imm_none) \
    X(OP_I32_NE, imm_none) \
    X(OP_I32_LT_S, imm_none) \
    X(OP_I32_LT_U, imm_none) \
    X(OP_I32_GT_S, imm_none) \
    X(OP_I32_GT_U, imm_none) \
    X(OP_I32_LE_S, imm_none) \
    X(OP_I32_LE_U, imm_none) \
    X(OP_I32_GE_S, imm_none) \
    X(OP_I32_GE_U, imm_none) \
    X(OP_I64_EQZ, imm_none) \
    X(OP_I64_EQ, imm_none) \
    X(OP_I64_NE, imm_none) \
    X(OP_I64_LT_S, imm_none) \
    X(OP_I64_LT_U, imm_none) \
    X(OP_I64_GT_S, imm_none) \
    X(OP_I64_GT_U, imm_none) \
    X(OP_I64_LE_S, imm_none) \
    X(OP_I64_LE_U, imm_none) \
    X(OP_I64_GE_S, imm_none) \
    X(OP_I64_GE_U, imm_none) \
    X(OP_F32_EQ, imm_none) \
    X(OP_F32_NE, imm_none) \
    X(OP_F32_LT, imm_none) \
    X(OP_F32_GT, imm_none) \
    X(OP_F32_LE, imm_none) \
    X(OP_F32_GE, imm_none) \
    X(OP_F64_EQ, imm_none) \
    X(OP_F64_NE, imm_none) \
    X(OP_F64_LT, imm_none) \
    X(OP_F64_GT, imm_none) \
    X(OP_F64_LE, imm_none) \
    X(OP_F64_GE, imm_none) \
    X(OP_I32_CLZ, imm_none) \
    X(OP_I32_CTZ, imm_none) \
    X(OP_I32_POPCNT, imm_none) \
    X(OP_I32_ADD, imm_none) \
    X(OP_I32_SUB, imm_none) \
    X(OP_I32_MUL, imm_none) \
    X(OP_I32_DIV_S, imm_none) \
    X(OP_I32_DIV_U, imm_none) \
    X(OP_I32_REM_S, imm_none) \
    X(OP_I32_REM_U, imm_none) \
    X(OP_I32_AND, imm_none) \
    X(OP_I32_OR, imm_none) \
    X(OP_I32_XOR, imm_none) \
    X(OP_I32_SHL, imm_none) \
    X(OP_I32_SHR_S, imm_none) \
    X(OP_I32_SHR_U, imm_none) \
    X(OP_I32_ROTL, imm_none) \
    X(OP_I32_ROTR, imm_none) \
    X(OP_I64_CLZ, imm_none) \
    X(OP_I64_CTZ, imm_none) \
    X(OP_I64_POPCNT, imm_none) \
    X(OP_I64_ADD, imm_none) \
    X(OP_I64_SUB, imm_none) \
    X(OP_I64_MUL, imm_none) \
    X(OP_I64_DIV_S, imm_none) \
    X(OP_I64_DIV_U, imm_none) \
    X(OP_I64_REM_S, imm_none) \
    X(OP_I64_REM_U, imm_none) \
    X(OP_I64_AND, imm_none) \
    X(OP_I64_OR, imm_none) \
    X(OP_I64_XOR, imm_none) \
    X(OP_I64_SHL, imm_none) \
    X(OP_I64_SHR_S, imm_none) \
    X(OP_I64_SHR_U, imm_none) \
    X(OP_I64_ROTL, imm_none) \
    X(OP_I64_ROTR, imm_none) \
    X(OP_F32_ABS, imm_none) \
    X(OP_F32_NEG, imm_none) \
    X(OP_F32_CEIL, imm_none) \
    X(OP_F32_FLOOR, imm_none) \
    X(OP_F32_TRUNC, imm_none) \
    X(OP_F32_NEAREST, imm_none) \
    X(OP_F32_SQRT, imm_none) \
    X(OP_F32_ADD, imm_none) \
    X(OP_F32_SUB, imm_none) \
    X(OP_F32_MUL, imm_none) \
    X(OP_F32_DIV, imm_none) \
    X(OP_F32_MIN, imm_none) \
    X(OP_F32_MAX, imm_none) \
    X(OP_F32_COPYSIGN, imm_none) \
    X(OP_F64_ABS, imm_none) \
    X(OP_F64_NEG, imm_none) \
    X(OP_F64_CEIL, imm_none) \
    X(OP_F64_FLOOR, imm_none) \
    X(OP_F64_TRUNC, imm_none) \
    X(OP_F64_NEAREST, imm_none) \
    X(OP_F64_SQRT, imm_none) \
    X(OP_F64_ADD, imm_none) \
    X(OP_F64_SUB, imm_none) \
    X(OP_F64_MUL, imm_none) \
    X(OP_F64_DIV, imm_none) \
    X(OP_F64_MIN, imm_none) \
    X(OP_F64_MAX, imm_none) \
    X(OP_F64_COPYSIGN, imm_none) \
    X(OP_I32_WRAP_I64, imm_none) \
    X(OP_I32_TRUNC_S_F32, imm_none) \
    X(OP_I32_TRUNC_U_F32, imm_none) \
    X(OP_I32_TRUNC_S_F64, imm_none) \
    X(OP_I32_TRUNC_U_F64, imm_none) \
    X(OP_I64_EXTEND_S_I32, imm_none) \
    X(OP_I64_EXTEND_U_I32, imm_none) \
    X(OP_I64_TRUNC_S_F32, imm_none) \
    X(OP_I64_TRUNC_U_F32, imm_none) \
    X(OP_I64_TRUNC_S_F64, imm_none) \
    X(OP_I64_TRUNC_U_F64, imm_none) \
    X(OP_F32_CONVERT_S_I32, imm_none) \
    X(OP_F32_CONVERT_U_I32, imm_none) \
    X(OP_F32_CONVERT_S_I64, imm_none) \
    X(OP_F32_CONVERT_U_I64, imm_none) \
    X(OP_F32_DEMOTE_F64, imm_none) \
    X(OP_F64_CONVERT_S_I32, imm_none) \
    X(OP_F64_CONVERT_U_I32, imm_none) \
    X(OP_F64_CONVERT_S_I64, imm_none) \
    X(OP_F64_CONVERT_U_I64, imm_none) \
    X(OP_F64_PROMOTE_F32, imm_none) \
    X(OP_I32_REINTERPRET_F32, imm_none) \
    X(OP_I64_REINTERPRET_F64, imm_none) \
    X(OP_F32_REINTERPRET_I32, imm_none) \
    X(OP_F64_REINTERPRET_I64, imm_none) \
    X(OP_EXTRA, imm_extra)

static int (*const table[])(FILE *, struct imm *) =
{
#define X(x, y) [x] = y,
    OPS
#undef X
};

static const char *const op_str[] =
{
#define X(x, y) [x] = #x,
    OPS
#undef X
};

static bool scoped(const uint8_t op)
{
    switch (op)
    {
        case OP_BLOCK:
        case OP_LOOP:
        case OP_IF:
            return true;

        default:
            break;
    }

    return false;
}

static bool branch(const uint8_t op)
{
    switch (op)
    {
        case OP_BR:
        case OP_BR_IF:
            return true;

        default:
            break;
    }

    return false;
}

static bool cbranch(const uint8_t op)
{
    switch (op)
    {
        case OP_IF:
        case OP_ELSE:
            return true;

        default:
            break;
    }

    return false;
}

struct block
{
    uint8_t op;
    unsigned level;
    long pc;
    struct block *next, *prev;
};

struct blocks
{
    struct block *head, *tail;
};

static int append_lo(const struct lo *const lo, struct parse *const p)
{
    const size_t n = p->n_lo + 1;
    struct lo *const o = realloc(p->lo, n * sizeof *p->lo);

    if (!o)
    {
        fprintf(stderr, "%s: realloc(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    o[p->n_lo++] = *lo;
    p->lo = o;
    return 0;
}

static int add_block(struct blocks *const b, const uint8_t op,
    const unsigned level, FILE *const f)
{
    const long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__,
            strerror(errno));
        return -1;
    }

    struct block *const v = malloc(sizeof *v);

    if (!v)
    {
        fprintf(stderr, "%s: realloc(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    *v = (const struct block)
    {
        .op = op,
        .level = level,
        .pc = offset,
        .prev = b->tail
    };

    if (!b->head)
        b->head = v;
    else
        b->tail->next = v;

    b->tail = v;
    return 0;
}

static int add_branch(const uint8_t op, struct lo *const lo, const struct blocks *const b,
    const varuint32 relative_depth, const unsigned level, FILE *const f)
{
    if (relative_depth >= level)
    {
        fprintf(stderr, "%s: invalid relative_depth=%lu, current level=%u\n",
            __func__, (unsigned long)relative_depth, level);
        return -1;
    }

    const long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__,
            strerror(errno));
        return -1;
    }

    const struct block *tgt = NULL;

    for (const struct block *bl = b->tail; bl; bl = bl->prev)
        if (bl->level == level - relative_depth)
        {
            tgt = bl;
            break;
        }

    if (!tgt)
    {
        fprintf(stderr, "%s: could not find matching block, pc=%ld\n",
            __func__, offset);
        return -1;
    }

    const size_t n = lo->n + 1;
    struct lo_entry *const ve = realloc(lo->entries, n * sizeof *lo->entries);

    if (!ve)
    {
        fprintf(stderr, "%s: realloc(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    struct lo_entry *const e = &ve[lo->n++];

    *e = (const struct lo_entry){.op = op, .b = tgt, .pc = offset};

    if (tgt->op == OP_LOOP)
        e->dst = tgt->pc;

    lo->entries = ve;
    return 0;
}

static int add_cbranch(const uint8_t op, struct lo *const lo,
    const struct blocks *const b, const unsigned level, FILE *const f)
{
    const long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__,
            strerror(errno));
        return -1;
    }

    const struct block *tgt = NULL;

    for (const struct block *bl = b->tail; bl; bl = bl->prev)
        if (bl->op == OP_IF && bl->level == level)
        {
            tgt = bl;
            break;
        }

    if (!tgt)
    {
        fprintf(stderr, "%s: could not find matching block, pc=%ld\n",
            __func__, offset);
        return -1;
    }

    const size_t n = lo->n + 1;
    struct lo_entry *const ve = realloc(lo->entries, n * sizeof *lo->entries);

    if (!ve)
    {
        fprintf(stderr, "%s: realloc(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    struct lo_entry *const e = &ve[lo->n++];

    *e = (const struct lo_entry){.op = op, .b = tgt, .pc = offset};
    lo->entries = ve;
    return 0;
}

static int add_brtable(struct lo *const lo, const struct blocks *const b,
    const struct imm_table *const t, const unsigned level, FILE *const f)
{
    for (varuint32 i = 0; i < (t->target_count + 1); i++)
        if (add_branch(OP_BR_TABLE, lo, b, t->targets[i], level, f))
        {
            fprintf(stderr, "%s: add_branch failed\n", __func__);
            return -1;
        }

    return 0;
}

static int set_end(struct lo *const lo, const struct blocks *const b,
    const unsigned level, FILE *const f)
{
    long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    offset -= sizeof b->head->op;

    const struct block *tgt = NULL;

    for (const struct block *bl = b->tail; bl; bl = bl->prev)
        if (bl->level == level)
        {
            tgt = bl;
            break;
        }

    if (!tgt)
    {
        fprintf(stderr, "%s: could not find matching block for end, pc=%ld\n",
            __func__, offset);
        return -1;
    }

    for (size_t i = 0; i < lo->n; i++)
    {
        struct lo_entry *const e = &lo->entries[i];

        if (e->b == tgt)
        {
            if (tgt->op == OP_BLOCK)
                e->dst = offset;
            else if (tgt->op == OP_IF &&
                (e->op == OP_ELSE || (e->op == OP_IF && !e->dst)))
                e->dst = offset;
        }
    }

    return 0;
}

static int set_else(struct lo *const lo, const struct blocks *const b,
    const unsigned level, FILE *const f)
{
    long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__, strerror(errno));
        return -1;
    }

    const struct block *tgt = NULL;

    for (const struct block *bl = b->tail; bl; bl = bl->prev)
        if (bl->level == level && bl->op == OP_IF)
        {
            tgt = bl;
            break;
        }

    if (!tgt)
    {
        fprintf(stderr, "%s: could not find matching block for end, pc=%ld\n",
            __func__, offset);
        return -1;
    }

    for (size_t i = 0; i < lo->n; i++)
    {
        struct lo_entry *const e = &lo->entries[i];

        if (e->b == tgt && e->op == OP_IF)
            e->dst = offset;
    }

    return 0;
}

static void free_blocks(struct blocks *const b)
{
    for (struct block *bl = b->head; bl;)
    {
        struct block *next = bl->next;

        free(bl);
        bl = next;
    }
}

static int read_one(struct lo *const lo, struct blocks *const b,
    unsigned *const level, FILE *const f)
{
    int ret = -1;
    size_t n;
    uint8_t op;
    struct imm imm = {0};
    const long offset = ftell(f);

    if (offset < 0)
    {
        fprintf(stderr, "%s: ftell(3): %s\n", __func__, strerror(errno));
        goto end;
    }
    else if (!(n = fread(&op, sizeof op, 1, f)))
    {
        fprintf(stderr, "%s: fread(3) failed, ferror(3)=%d, feof(3)=%d\n",
            __func__, ferror(f), feof(f));
        goto end;
    }
    else if (op >= sizeof table / sizeof *table || !table[op])
    {
        fprintf(stderr, "%s: invalid or unsupported opcode %#" PRIx8 ", "
            "offset=%#lx\n", __func__, op, offset);
        goto end;
    }
    else if (table[op](f, &imm))
    {
        fprintf(stderr, "%s: failed to read opcode %#" PRIx8 "\n",
            __func__, op);
        goto end;
    }
    else if (scoped(op) && add_block(b, op, ++(*level), f))
    {
        fprintf(stderr, "%s: add_block failed\n", __func__);
        goto end;
    }
    else if (branch(op) && add_branch(op, lo, b, imm.u.varuint32, *level, f))
    {
        fprintf(stderr, "%s: add_branch failed\n", __func__);
        goto end;
    }
    else if (cbranch(op) && add_cbranch(op, lo, b, *level, f))
    {
        fprintf(stderr, "%s: add_cbranch failed\n", __func__);
        goto end;
    }
    else if (op == OP_BR_TABLE && add_brtable(lo, b, &imm.u.table, *level, f))
    {
        fprintf(stderr, "%s: add_brtable failed\n", __func__);
        goto end;
    }
    else if (op == OP_ELSE && set_else(lo, b, *level, f))
    {
        fprintf(stderr, "%s: set_else failed\n", __func__);
        goto end;
    }
    else if (op == OP_END)
    {
        if (*level)
        {
            if (set_end(lo, b, *level, f))
            {
                fprintf(stderr, "%s: set_end failed\n", __func__);
                goto end;
            }

            (*level)--;
        }
        else
        {
            ret = 1;
            goto end;
        }
    }

    ret = 0;

end:
    free_imm(&imm);
    return ret;
}

int read_instr(FILE *const f, struct parse *const p)
{
    unsigned level = 0;
    struct lo lo = {0};
    struct blocks b = {0};

    for (;;)
    {
        const int n = read_one(&lo, &b, &level, f);

        if (n < 0)
            goto failure;
        else if (n)
            break;
    }

    if (append_lo(&lo, p))
    {
        fprintf(stderr, "%s: append_lo failed\n", __func__);
        goto failure;
    }

    free_blocks(&b);
    return 0;

failure:
    free(lo.entries);
    free_blocks(&b);
    return -1;
}