/*
 * nanowasm, a tiny WebAssembly/Wasm interpreter
 * Copyright (C) 2023-2025  Xavier Del Campo Romero
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include <nanowasm/nw.h>
#include <errno.h>
#include <stddef.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct inst
{
    long retval;

    struct mem
    {
        void *p;
        size_t n;
    } stack, linear, global;
};

static int io_read(void *const buf, const size_t n, void *const user)
{
    FILE *const f = user;
    const size_t r = fread(buf, 1, 1, f);

    if (!r && ferror(f))
    {
        fprintf(stderr, "%s: ferror\n", __func__);
        return -1;
    }

    return r;
}

static enum nw_state io_seek(const long offset, void *const user)
{
    if (fseek(user, offset, SEEK_SET))
    {
        fprintf(stderr, "%s: fseek(3): %s\n", __func__, strerror(errno));
        return NW_FATAL;
    }

    return NW_OK;
}

static enum nw_state io_tell(long *const out, void *const user)
{
    const long offset = ftell(user);

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

    *out = offset;
    return NW_OK;
}

static int io_eof(void *const user)
{
    FILE *const f = user;

    return feof(f);
}

static int push(const void *const src, const size_t n, void *const user)
{
    struct inst *const inst = user;
    struct mem *const s = &inst->stack;
    char *const p = realloc(s->p, s->n + n);

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

    fprintf(stderr, "pushed %zu bytes: {", n);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)src)[i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    memcpy(&p[s->n], src, n);
    s->p = p;
    s->n += n;
    fprintf(stderr, "}, stack size=%zu\n", s->n);
    return n;
}

static int pop(void *const dst, const size_t n, void *const user)
{
    struct inst *const inst = user;
    struct mem *const s = &inst->stack;

    if (s->n < n)
    {
        fprintf(stderr, "%s: stack underflow\n", __func__);
        return -1;
    }

    const char *const src = (const char *)s->p + (s->n - n);

    fprintf(stderr, "popped %zu bytes: {", n);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)src)[i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    memcpy(dst, src, n);

    if (s->n == n)
    {
        free(s->p);
        s->p = NULL;
    }
    else
    {
        char *const p = realloc(s->p, s->n - n);

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

        s->p = p;
    }

    s->n -= n;
    fprintf(stderr, "}, stack size=%zu\n", s->n);
    return n;
}

static size_t ptr(void *const user)
{
    return ((const struct inst *)user)->stack.n;
}

static int load(const nw_varuint32 offset, void *const dst, const size_t n,
    const struct mem *const m)
{
    if (n > m->n || offset > m->n - n)
    {
        fprintf(stderr, "%s: out-of-bounds access, offset=%lu, n=%zu\n",
            __func__, (unsigned long)offset, n);
        return -1;
    }

    fprintf(stderr, "%s: loaded %zu bytes from offset %lu: {", __func__, n,
        (unsigned long)offset);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)m->p)[offset + i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    fputs("}\n", stderr);
    memcpy(dst, (const char *)m->p + offset, n);
    return n;
}

static int store(const nw_varuint32 offset, const void *const src,
    const size_t n, const struct mem *const m)
{
    if (n > m->n || offset > m->n - n)
    {
        fprintf(stderr, "%s: out-of-bounds access, offset=%lu, n=%zu\n",
            __func__, (unsigned long)offset, n);
        return -1;
    }

    fprintf(stderr, "%s: stored %zu bytes into offset %lu: {", __func__, n,
        (unsigned long)offset);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)src)[i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    fputs("}\n", stderr);
    memcpy((char *)m->p + offset, src, n);
    return n;
}

static int stread(const size_t offset, void *const dst, const size_t n,
    void *const user)
{
    const struct inst *const inst = user;

    return load(offset, dst, n, &inst->stack);
}

static int stwrite(const size_t offset, const void *const src, const size_t n,
    void *const user)
{
    const struct inst *const inst = user;

    return store(offset, src, n, &inst->stack);
}

static int ensure_linear(struct mem *const m, const unsigned long offset,
    const size_t n)
{
    void *const p = realloc(m->p, offset + n);

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

    m->p = p;
    m->n = offset + n;
    memset((char *)m->p + offset, 0, n);
    return 0;
}

static int lnload(const unsigned long offset, void *const dst, const size_t n,
    void *const user)
{
    struct inst *const i = user;
    struct mem *const m = &i->linear;

    if (offset >= m->n && ensure_linear(m, offset, n))
    {
        fprintf(stderr, "%s: ensure_linear failed\n", __func__);
        return -1;
    }

    fprintf(stderr, "%s: loaded %zu bytes from offset %lu: {", __func__, n,
        (unsigned long)offset);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)m->p)[offset + i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    fputs("}\n", stderr);
    memcpy(dst, (const char *)m->p + offset, n);
    return n;
}

static int lnstore(const unsigned long offset, const void *const src,
    const size_t n, void *const user)
{
    struct inst *const i = user;
    struct mem *const m = &i->linear;

    if (offset >= m->n && ensure_linear(m, offset, n))
    {
        fprintf(stderr, "%s: ensure_linear failed\n", __func__);
        return -1;
    }

    fprintf(stderr, "%s: stored %zu bytes into offset %lu: {", __func__, n,
        (unsigned long)offset);

    for (size_t i = 0; i < n; i++)
    {
        fprintf(stderr, "%hhu", ((const char *)src)[i]);

        if (i + 1 < n)
            fputs(", ", stderr);
    }

    fputs("}\n", stderr);
    memcpy((char *)m->p + offset, src, n);
    return n;
}

static int glload(const unsigned long offset, void *const dst,
    const size_t n, void *const user)
{
    const struct inst *const i = user;

    return load(offset, dst, n, &i->global);
}

static int glstore(const unsigned long offset, const void *const src,
    const size_t n, void *const user)
{
    const struct inst *const i = user;

    return store(offset, src, n, &i->global);
}

static enum nw_state wasi_exit(const union nw_value *const params,
    union nw_value *const ret, void *user, struct nw_next *const next)
{
    struct inst *const i = user;

    i->retval = params->i32;
    return NW_OK;
}

int main(int argc, char *argv[])
{
    int ret = EXIT_FAILURE;
    FILE *f = NULL;
    struct inst ins = {0};

    if (argc != 2)
    {
        fprintf(stderr, "%s <wasm>\n", *argv);
        goto end;
    }

    const char *const path = argv[1];

    if (!(f = fopen(path, "rb")))
    {
        fprintf(stderr, "%s: fopen(3) %s: %s\n", __func__, path,
            strerror(errno));
        goto end;
    }

    static const struct nw_import imports[] =
    {
        {
            .kind = NW_KIND_FUNCTION,
            .module = "wasi_snapshot_preview1",
            .field = "proc_exit",
            .u.function =
            {
                .fn = wasi_exit,
                .signature = "(i)"
            }
        }
    };

    const struct nw_io_cfg io =
    {
        .read = io_read,
        .seek = io_seek,
        .tell = io_tell,
        .eof = io_eof,
        .user = f
    };

    enum {N_IMPORTS = sizeof imports / sizeof *imports};

    const struct nw_mod_cfg cfg =
    {
        .imports = imports,
        .imp_indexes = (struct nw_import_index[N_IMPORTS]){{0}},
        .n_imports = N_IMPORTS,
        .io = io
    };

    struct nw_mod m;
    struct nw_mod_out mout;

    nw_init(&m, &cfg);

again:

    switch (nw_load(&m, &mout))
    {
        case NW_OK:
            break;

        case NW_AGAIN:
            goto again;

        case NW_FATAL:
            fprintf(stderr, "%s: nw_load failed\n", __func__);
            goto end;
    }

    if (!(ins.global.p = malloc(ins.global.n = mout.global)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        goto end;
    }

    enum {ARGS = 1};
    const struct nw_inst_cfg icfg =
    {
        .entry = "_start",
        .interp_cfg =
        {
            .io = io,
            .m = &m,
            .user = &ins,
            .args = (union nw_value[ARGS]){{0}},
            .n_args = ARGS,
            .stack =
            {
                .push = push,
                .pop = pop,
                .ptr = ptr,
                .read = stread,
                .write = stwrite
            },

            .linear =
            {
                .load = lnload,
                .store = lnstore
            },

            .global =
            {
                .load = glload,
                .store = glstore
            }
        }
    };

    struct nw_inst inst;

    if (nw_start(&inst, &icfg))
    {
        fprintf(stderr, "%s: nw_start failed\n", __func__);
        goto end;
    }

again2:

    switch (nw_run(&inst))
    {
        case NW_OK:
            break;

        case NW_AGAIN:

            if (ins.retval)
            {
                ret = ins.retval;
                fprintf(stderr, "instance exited with status %d\n", ret);
                goto end;
            }

            goto again2;

        case NW_FATAL:
            fprintf(stderr, "%s: nw_run failed: %s\n", __func__,
                nw_rexc(&inst));
            goto end;
    }

    ret = EXIT_SUCCESS;

end:

    if (f && fclose(f))
    {
        fprintf(stderr, "%s: fclose(3) %s: %s\n", __func__, path,
            strerror(errno));
        goto end;
    }

    free(ins.stack.p);
    free(ins.linear.p);
    free(ins.global.p);
    return ret;
}