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

static enum nw_state seek_param_types(struct nw_interp *);

static enum nw_state seek_pc(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = cfg->seek(ci->pc, cfg->user);

    if (n)
        return n;

    nwp_interp_resume(i);
    return NW_AGAIN;
}

static enum nw_state push_return(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const enum nw_state n = nwp_stack_push(i, &ci->io);

    if (n)
        return n;

    i->next = seek_pc;
    return NW_AGAIN;
}

static int parse_signature(const struct nw_import *const imp,
    size_t *const nargs, int *const ret)
{
    enum
    {
        RETURN,
        PARAM_BEGIN,
        PARAM,
        END
    } state = 0;

    const char *s;

    if (nargs)
        *nargs = 0;

    if (ret)
        *ret = 0;

    for (s = imp->u.function.signature; *s; s++)
    {
        switch (state)
        {
            case RETURN:
                if (strchr("iIfF", *s))
                {
                    if (ret)
                        *ret = 1;

                    state++;
                }
                else if (*s == '(')
                    state = PARAM;

                break;

            case PARAM_BEGIN:
                if (*s != '(')
                    return -1;

                state++;
                break;

            case PARAM:
                if (strchr("iIfF", *s))
                {
                    if (nargs)
                        (*nargs)++;
                }
                else if (*s == ')')
                    state++;
                else
                    return -1;

                break;

            case END:
                goto end;
        }
    }

end:
    return state != END || *s;
}

static int check_ret(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const struct nw_import *const imp = ci->imp;
    const struct nw_fn *const fn = &ci->fn;
    const nw_varuint1 retcnt = fn->ret.count;
    int hasret;

    if (parse_signature(imp, NULL, &hasret))
    {
        static const char *const exc = "invalid imported function signature";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s::%s: %s: %s\n", imp->module, imp->field, exc,
            imp->u.function.signature);
#endif
        return NW_FATAL;
    }
    else if (hasret != retcnt)
    {
        static const char *const exc = "return count mismatch";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s::%s: %s: expected: %d, got: %d\n", imp->module, imp->field,
            exc, hasret, retcnt);
#endif
        return -1;
    }

    return 0;
}

static enum nw_state call_done(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const struct nw_fn *const fn = &ci->fn;
    const nw_varuint32 pcnt = fn->param_count;
    const struct nw_interp_cfg *const icfg = &i->cfg;
    union nw_value *const ret = fn->ret.count ? icfg->args + pcnt : NULL;
    struct nw_return r = {0};

    if (ret)
    {
        size_t sz;
        const enum nw_type type = fn->ret.type;

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

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

            io.buf = ret;
            io.n = sz;
            ci->io = io;
            i->next = push_return;
        }

        r.count = 1;
        r.type = type;
    }
    else
        nwp_interp_resume(i);

    i->fr.prev_ret = r;
    return NW_AGAIN;
}

static enum nw_state repeat(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    struct nw_next *const next = &ci->next;
    const struct nw_import *const imp = ci->imp;
    const enum nw_state n = next->fn(next->user, next);

    if (n == NW_FATAL)
    {
        static const char *const exc = "import function failed";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %s::%s\n", exc, imp->module, imp->field);
#endif
        return NW_FATAL;
    }
    else if (n)
        return n;

    return call_done(i);
}

static enum nw_state call(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const struct nw_fn *const fn = &ci->fn;
    const nw_varuint32 pcnt = fn->param_count;
    const struct nw_interp_cfg *const icfg = &i->cfg;
    const union nw_value *const args = pcnt ? icfg->args : NULL;
    union nw_value *const ret = fn->ret.count ? icfg->args + pcnt : NULL;
    const struct nw_import *const imp = ci->imp;
    enum nw_state n;

    if (fn->ret.count && check_ret(i))
        return NW_FATAL;
    else if ((n = imp->u.function.fn(args, ret, icfg->user, &ci->next)))
    {
        if (n == NW_FATAL)
        {
            static const char *const exc = "import function failed";

            i->exception = exc;
#ifdef NW_LOG
            nwp_log("%s: %s::%s\n", exc, imp->module, imp->field);
#endif
        }
        else if (ci->next.fn)
            i->next = repeat;

        return n;
    }

    return call_done(i);
}

static enum nw_state pop_args(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const enum nw_state n = nwp_stack_pop(i, &ci->io);

    if (n)
        return n;

    i->next = call;
    return NW_AGAIN;
}

static enum nw_state get_value(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const enum nw_state n = nwp_stack_read(i, &ci->io, ci->addr);
    const struct nw_import *const imp = ci->imp;
    const size_t max = i->cfg.n_args;
    const nw_varuint32 pcnt = ci->fn.param_count;
    size_t nargs;

    if (n)
        return n;
    else if (parse_signature(imp, &nargs, NULL))
    {
        static const char *const exc = "invalid imported function signature";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s::%s: %s: %s\n", imp->module, imp->field, exc,
            imp->u.function.signature);
#endif
        return NW_FATAL;
    }
    else if (nargs > max)
    {
        static const char *const exc = "exceeded maximum args";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s, %s::%s requests %lu, max: %lu\n", exc, imp->module,
            imp->field, (unsigned long)nargs, (unsigned long)max);
#endif
        return NW_FATAL;
    }
    else if (pcnt != nargs)
    {
        static const char *const exc = "arg number mismatch";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s, %s::%s expected %lu args, got %lu\n", exc, imp->module,
            imp->field, (unsigned long)nargs, (unsigned long)pcnt);
#endif
        return NW_FATAL;
    }

    i->cfg.args[ci->param_i++] = ci->value;

    if (ci->param_i >= ci->fn.param_count)
    {
        struct nw_sm_io io = {0};

        io.n = ci->sz;
        ci->io = io;
        i->next = pop_args;
    }
    else
        i->next =  seek_param_types;

    return NW_AGAIN;
}

static enum nw_state get_param_type(struct nw_interp *const i)
{
    nw_varint7 type;
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    struct nw_sm_leb128 *const l = &ci->leb128;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = nwp_varint7(cfg, l, &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;
    }
    else if (ci->type_i >= ci->param_i)
    {
        if (ci->addr < sz)
        {
            static const char *const exc = "stack underflow";

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

        if (ci->type_i == ci->param_i)
            ci->type = vtype;

        ci->addr -= sz;
    }

    ci->sz += sz;

    if (++ci->type_i >= ci->fn.param_count)
    {
        struct nw_sm_io io = {0};

        if (nwp_type_sz(ci->type, &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;
        }

        io.buf = &ci->value;
        io.n = sz;
        ci->io = io;
        i->next = get_value;
    }

    return NW_AGAIN;
}

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

    if (n)
        return n;

    ci->type_i = 0;
    ci->sz = 0;
    ci->addr = nwp_stack_ptr(i);
    i->next = get_param_type;
    return NW_AGAIN;
}

static enum nw_state tell(struct nw_interp *const i)
{
    struct nw_i_sm_call_import *const ci = &i->sm.call_import;
    const struct nw_io_cfg *const cfg = &i->cfg.io;
    const enum nw_state n = cfg->tell(&ci->pc, cfg->user);

    if (n)
        return n;

    i->next = ci->fn.param_count ? seek_param_types : call;
    return NW_AGAIN;
}

static const struct nw_import *find(const struct nw_interp *const i,
    const nw_varuint32 index)
{
    const struct nw_mod_cfg *const cfg = &i->cfg.m->cfg;
    size_t j;

    for (j = 0; j < cfg->n_imports; j++)
    {
        const struct nw_import_index *const ii = &cfg->imp_indexes[j];

        if (ii->index == index)
        {
            const struct nw_import *const imp = &cfg->imports[j];

#ifdef NW_LOG
            nwp_log("import function index %lu matches %s::%s (import index %lu)\n",
                (unsigned long)index, imp->module, imp->field, (unsigned long)j);
#endif
            return &cfg->imports[j];
        }
    }

    return NULL;
}

static enum nw_state prepare(struct nw_interp *const i)
{
    const struct nw_fn fn = i->sm.type.out;
    const struct nw_import *const imp = find(i, fn.index);

    if (!imp)
    {
        static const char *const exc = "function import not found";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %lu\n", exc, (unsigned long)fn.index);
#endif
        return NW_FATAL;
    }
    else if (imp->kind != NW_KIND_FUNCTION)
    {
        static const char *const exc = "import not a function";

        i->exception = exc;
#ifdef NW_LOG
        nwp_log("%s: %s::%s\n", exc, imp->module, imp->field);
#endif
        return NW_FATAL;
    }
    else
    {
        const struct nw_i_sm_call_import ci = {0};
        struct nw_i_sm_call_import *const pci = &i->sm.call_import;

        *pci = ci;
        pci->fn = fn;
        pci->imp = imp;
        i->next = tell;
    }

    return NW_AGAIN;
}

void nwp_call_import(struct nw_interp *const i, const nw_varuint32 index)
{
    nwp_get_function_type(i, index, prepare);
}