/*
 * 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 <nw/routines.h>
#include <nw/interp.h>
#include <nw/io.h>
#include <nw/log.h>
#include <nw/stack.h>
#include <nw/types.h>

static enum nw_state prepare_pop(struct nw_interp *);

static enum nw_state seek_pc(struct nw_interp *const i)
{
    const struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = cfg->seek(u->fr.pc, cfg->user);
    struct nw_return r;

    if (n)
        return n;

#ifdef NW_LOG
    nwp_log("returning to function %lu\n", u->fr.fn.index);
#endif

    r = i->fr.fn.ret;
    i->fr = u->fr;
    i->fr.prev_ret = r;
    nwp_interp_resume(i);
    return NW_AGAIN;
}

static enum nw_state push_return_value(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const enum nw_state n = nwp_stack_push(i, &u->io);

    if (n)
        return n;
    else if (u->fr.child)
        i->next = seek_pc;
    else
        return NW_OK;

    return NW_AGAIN;
}

static int prepare_retval(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_fn *const fn = &i->fr.fn;
    const struct nw_return *const r = &fn->ret;
    const enum nw_type t = r->type;
    size_t sz;

    if (nwp_type_sz(t, &sz))
    {
        static const char *const exc = "invalid return type";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %#x\n", exc, (unsigned)t);
#endif
        return -1;
    }
    else
    {
        struct nw_sm_io io = {0};

        io.buf = &u->retval;
        io.n = sz;
        u->io = io;
        i->push_type = t;
        i->next = push_return_value;
    }

    return 0;
}

static enum nw_state pop_params(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_fn *const fn = &i->fr.fn;
    const struct nw_return *const r = &fn->ret;
    const enum nw_state n = nwp_stack_pop(i, &u->io);

    if (n)
        return n;
    else if (r->count)
    {
        if (prepare_retval(i))
            return NW_FATAL;
    }
    else if (u->fr.child)
        i->next = seek_pc;
    else
        return NW_OK;

    return NW_AGAIN;
}

static enum nw_state get_param_type(struct nw_interp *const i)
{
    nw_varint7 type;
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = nwp_varint7(cfg, &u->leb128, &type, cfg->user);
    enum nw_type vtype;
    size_t sz;

    if (n)
        return n;
    else if (nwp_get_type(type, &vtype) || nwp_type_sz(vtype, &sz))
    {
        static const char *const exc = "invalid param type";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %#x\n", exc, (unsigned)type);
#endif
        return NW_FATAL;
    }

    u->sz += sz;

    if (++u->entry_i >= i->fr.fn.param_count)
    {
        struct nw_sm_io io = {0};

        io.n = u->sz;
        u->io = io;
        i->next = pop_params;
    }

    return NW_AGAIN;
}

static enum nw_state seek_param_types(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = cfg->seek(i->fr.fn.param_types, cfg->user);

    if (n)
        return n;

    u->entry_i = 0;
    i->next = get_param_type;
    return NW_AGAIN;
}

static enum nw_state check_retval(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_fn *const fn = &i->fr.fn;
    const size_t addr = nwp_stack_ptr(i), start = i->fr.local_start;
    const struct nw_return *const r = &fn->ret;

    if (addr != start)
    {
        static const char *const exc = "mismatched stack address";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: expected: %lu, got: %lu\n", exc,
            (unsigned long)start, (unsigned long)addr);
#endif
        return NW_FATAL;
    }
    else if (fn->param_count)
        i->next = seek_param_types;
    else if (r->count)
    {
        if (prepare_retval(i))
            return NW_FATAL;
    }
    else if (u->fr.child)
        i->next = seek_pc;
    else
        return NW_OK;

    return NW_AGAIN;
}

static enum nw_state pop_meta(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const enum nw_state n = nwp_stack_pop(i, &u->io);

    if (n)
        return n;
    else if (u->pending)
        nwp_find_local(i, &u->fl, u->pending - 1, prepare_pop, NULL);
    else
        return check_retval(i);

    return NW_AGAIN;
}

static enum nw_state pop_local(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    struct nw_local_meta *const m = &u->fl.meta;
    const enum nw_state n = nwp_stack_pop(i, &u->io);

    if (n)
        return n;

    u->pending--;

    if (++u->entry_i >= m->entry_count)
    {
        struct nw_sm_io io = {0};

        io.buf = m;
        io.n = sizeof *m;
        u->io = io;
        u->entry_i = 0;
        i->next = pop_meta;
    }
    else if (u->pending)
        return prepare_pop(i);
    else
    {
        static const char *const exc = "entry and local count mismatch";
#ifdef NW_LOG
        const unsigned long pending = m->entry_count - u->entry_i;
#endif

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %lu locals pending\n", exc, pending);
#endif
        return NW_FATAL;
    }

    return NW_AGAIN;
}

static enum nw_state prepare_pop(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_local_meta *const m = &u->fl.meta;
    size_t sz;

    if (nwp_type_sz(m->type, &sz))
    {
        static const char *const exc = "invalid type";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %#x\n", exc, (unsigned)m->type);
#endif
        return NW_FATAL;
    }
    else
    {
        struct nw_sm_io io = {0};

        io.buf = &u->value;
        io.n = sz;
        u->io = io;
        i->next = pop_local;
    }

    return NW_AGAIN;
}

static enum nw_state pop_frame(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_frame *const fr = &i->fr;
    const enum nw_state n = nwp_stack_pop(i, &u->io);

    if (n)
        return n;
    else if (fr->local_count)
    {
        u->pending = fr->local_count;
        nwp_find_local(i, &u->fl, u->pending - 1, prepare_pop, NULL);
    }
    else
        return check_retval(i);

    return NW_AGAIN;
}

static enum nw_state pop_return_value(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const enum nw_state n = nwp_stack_pop(i, &u->io);

    if (n)
        return n;
    else
    {
        struct nw_sm_io io = {0};

        io.buf = &u->fr;
        io.n = sizeof u->fr;
        u->io = io;
        i->next = pop_frame;
    }

    return NW_AGAIN;
}

static enum nw_state run(struct nw_interp *const i)
{
    struct nw_i_sm_unwind *const u = &i->sm.unwind;
    const struct nw_frame *const fr = &i->fr;
    const struct nw_fn *const fn = &fr->fn;
    const struct nw_return *const r = &fn->ret;

    if (r->count)
    {
        size_t sz, end, addr;

        if (nwp_type_sz(r->type, &sz))
        {
            static const char *const exc = "invalid type";

            i->exception = exc;
#ifdef NW_LOG
            nwp_log("%s: %#x\n", exc, (unsigned)r->type);
#endif
            return NW_FATAL;
        }

        end = fr->local_end;
        addr = nwp_stack_ptr(i);

        if (addr < sz || addr < sizeof *fr)
        {
            static const char *const exc = "stack underflow";

            i->exception = exc;
#ifdef NW_LOG
            nwp_log("%s\n", exc);
#endif
            return NW_FATAL;
        }
        else if ((addr -= sizeof *fr + sz) != end)
        {
            static const char *const exc = "stack memory leak";

            i->exception = exc;
#ifdef NW_LOG
            nwp_log("%s: expected: %lu, got: %lu\n", exc,  (unsigned long)end,
            (unsigned long)addr);
#endif
            return NW_FATAL;
        }
        else
        {
            struct nw_sm_io io = {0};

            io.buf = &u->retval;
            io.n = sz;
            u->io = io;
            i->next = pop_return_value;
        }
    }
    else
    {
        struct nw_sm_io io = {0};

        io.buf = &u->fr;
        io.n = sizeof u->fr;
        u->io = io;
        i->next = pop_frame;
    }

    return NW_AGAIN;
}

void nwp_unwind(struct nw_interp *const i)
{
    const struct nw_i_sm_unwind u = {0};

    i->sm.unwind = u;
    i->next = run;
}