/*
 * wnix, a Unix-like operating system for WebAssembly applications.
 * 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 <bin.h>
#include <bin/proc.h>
#include <bin/types.h>
#include <aio.h>
#include <kprintf.h>
#include <loop.h>
#include <state.h>
#include <dynstr.h>
#include <nanowasm/nw.h>
#include <errno.h>
#include <limits.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

static enum nw_state check_start(void *, struct nw_next *);
static int read_string(struct bin_proc *, struct nw_next *);

static int read_done(const enum state state, void *const args)
{
    struct bin_proc *const p = args;
    struct bin_dbg *const d = &p->dbg;

    if (loop_rm_aio(d->aio) || bin_proc_unlock(p))
        return -1;

    aio_free(d->aio);
    d->aio = NULL;
    return 0;
}

static int read_byte(struct bin_proc *const p)
{
    struct bin_dbg *const d = &p->dbg;

    const struct fs_read r =
    {
        .fd = &d->fd,
        .buf = &d->b,
        .n = sizeof d->b
    };

    const struct aio_done done =
    {
        .f = read_done,
        .args = p
    };

    struct aio *const aio = aio_read(&r, &done);

    if (!aio
        || loop_add_aio(aio)
        || bin_proc_lock(p))
        goto failure;

    d->aio = aio;
    return 0;

failure:
    aio_free(aio);
    return -1;
}

static const char *token(const char *s, size_t *const n,
    const char **const next)
{
    const char *const sep = strchr(s, ':');

    if (sep)
    {
        *n = sep - s;
        *next = sep + 1;
    }
    else
    {
        *n = sep ? sep - s : strlen(s);
        *next = NULL;
    }

    return s;
}

static int parse_num(const char *const s, unsigned long *const out)
{
    char *end;
    unsigned long v;

    errno = 0;
    v = strtoul(s, &end, 0);

    if (errno || *end)
        return -1;

    *out = v;
    return 0;
}

static int replied(const enum state state, void *const args)
{
    struct bin_proc *const p = args;
    struct bin_dbg *const d = &p->dbg;

    if (loop_rm_aio(d->aio) || bin_proc_unlock(p))
        return -1;

    aio_free(d->aio);
    d->aio = NULL;
    return 0;
}

static enum nw_state finished(void *const user, struct nw_next *const next)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const d = &p->dbg;

    dynstr_free(&d->dstr);

    if (d->running)
        return NW_OK;

    *next = (const struct nw_next){.fn = check_start, .user = p};
    return NW_AGAIN;
}

static enum nw_state reply(struct bin_proc *const p, struct nw_next *const next,
    const char *const s)
{
    struct aio *aio = NULL;
    struct bin_dbg *const d = &p->dbg;
    const struct fs_write w =
    {
        .fd = &d->fd,
        .buf = s,
        .n = strlen(s)
    };

    const struct aio_done done =
    {
        .f = replied,
        .args = p
    };

    if (!(aio = aio_write(&w, &done))
        || loop_add_aio(aio)
        || bin_proc_lock(p))
        goto failure;

    d->aio = aio;
    *next = (const struct nw_next){.fn = finished, .user = p};
    return NW_AGAIN;

failure:
    aio_free(aio);
    return NW_FATAL;
}

static int append_value(struct dynstr *const d,
    const struct nw_dbg_value *const v)
{
    /* TODO: is %f valid for both float and double? */
    static const char *const fmt[] =
    {
        [NW_TYPE_I32] = "%lu",
        [NW_TYPE_I64] = "%llu",
        [NW_TYPE_F32] = "%f",
        [NW_TYPE_F64] = "%f"
    };

    return dynstr_append(d, fmt[v->type], v->value);
}

static enum nw_state write_value(void *const user, struct nw_next *const next)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const db = &p->dbg;
    struct dynstr *const d = &db->dstr;
    struct nw_dbg_value v;

    dynstr_free(d);

    if (nw_dbg_value(&db->dbg, &v)
        || dynstr_append(d, ";ok:")
        || append_value(d, &v)
        || dynstr_append(d, "\n"))
        return NW_FATAL;

    return reply(user, next, d->str);
}

static enum nw_state trap(void *const user, struct nw_next *const next)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const db = &p->dbg;
    struct dynstr *const d = &db->dstr;

    dynstr_init(d);

    if (dynstr_append(d, ";trap:%#lx\n", db->offset))
        return NW_FATAL;

    return reply(user, next, d->str);
}

static enum nw_state ok(void *const user, struct nw_next *const next)
{
    return reply(user, next, ";ok\n");
}

static enum nw_state nok(void *const user, struct nw_next *const next)
{
    return reply(user, next, ";nok\n");
}

static int add_bkpt(struct bin_dbg *const d, const long addr)
{
    const size_t n = d->n_bkpt + 1;
    long *const b = realloc(d->bkpt, n * sizeof *d->bkpt);

    if (!b)
        return -1;

    b[d->n_bkpt] = addr;
    d->bkpt = b;
    d->n_bkpt = n;
    return 0;
}

static int cmd_b(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    unsigned long addr;
    struct bin_dbg *const d = &p->dbg;
    const size_t n = d->n_bkpt;
    const char *ntok, *const tk = token(s, &tklen, &ntok);

    if (!tk || ntok)
    {
        kprintf("Syntax: b:<pc>\n");
        return 1;
    }
    else if (parse_num(tk, &addr) || addr > LONG_MAX)
    {
        kprintf("Invalid value: %s\n", tk);
        return 1;
    }
    else if (add_bkpt(d, addr))
        return -1;
    else
        kprintf("Set breakpoint %zu at address %#lx\n", n, addr);

    *next = (const struct nw_next){.fn = ok, .user = p};
    return 0;
}

static int cmd_bt(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    /* TODO */
    kprintf("%s: TODO\n", __func__);
    return 1;
}

static int cmd_c(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    struct bin_dbg *const d = &p->dbg;

    if (s)
    {
        kprintf("Syntax: c\n");
        return 1;
    }

    d->running = true;
    *next = (const struct nw_next){.fn = ok, .user = p};
    return 0;
}

static int resize_bkpt(struct bin_dbg *const d, const unsigned long index)
{
    const size_t n = d->n_bkpt - 1;
    long *const dst = &d->bkpt[index];

    if (index < n)
        memcpy(dst, dst + 1, (d->n_bkpt - index) * sizeof *dst);

    if (n)
    {
        long *const b = realloc(d->bkpt, n * sizeof *d->bkpt);

        if (!b)
            return -1;

        d->bkpt = b;
    }
    else
    {
        free(d->bkpt);
        d->bkpt = NULL;
    }

    d->n_bkpt = n;
    return 0;
}

static void delete_bkpt(struct bin_dbg *const d)
{
    free(d->bkpt);
    d->bkpt = NULL;
    d->n_bkpt = 0;
}

static int cmd_d(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    unsigned long index;
    struct bin_dbg *const d = &p->dbg;
    const char *ntok, *const tk = token(s, &tklen, &ntok);

    if (!tk)
    {
        if (d->n_bkpt)
        {
            delete_bkpt(d);
            kprintf("Deleted all breakpoints\n");
        }
        else
            kprintf("No breakpoints to delete\n");
    }
    else if (ntok)
    {
        kprintf("Syntax: d[:<b-index>]\n");
        return 1;
    }
    else if (parse_num(tk, &index) || index >= d->n_bkpt)
    {
        kprintf("Invalid breakpoint index: %s\n", tk);
        return 1;
    }
    else if (resize_bkpt(d, index))
        return -1;
    else
        kprintf("Deleted breakpoint index %lu\n", index);

    *next = (const struct nw_next){.fn = ok, .user = p};
    return 0;
}

static int cmd_g(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    unsigned long index;
    struct bin_dbg *const d = &p->dbg;
    const char *ntok, *const tk = token(s, &tklen, &ntok);

    if (!tk || ntok)
    {
        kprintf("Syntax: g:<index>\n");
        return 1;
    }
    else if (parse_num(tk, &index))
    {
        kprintf("Invalid value: %s\n", tk);
        return 1;
    }

    nw_dbg_global(&d->dbg, index);
    *next = (const struct nw_next){.fn = write_value, .user = p};
    return 0;
}

static int cmd_l(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    unsigned long index;
    struct bin_dbg *const d = &p->dbg;
    const char *ntok, *const tk = token(s, &tklen, &ntok);

    if (!tk || ntok)
    {
        kprintf("Syntax: l:<index>\n");
        return 1;
    }
    else if (parse_num(tk, &index))
    {
        kprintf("Invalid value: %s\n", tk);
        return 1;
    }

    nw_dbg_local(&d->dbg, index);
    *next = (const struct nw_next){.fn = write_value, .user = p};
    return 0;
}

static int cmd_p(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    unsigned long index;
    struct bin_dbg *const d = &p->dbg;
    const char *ntok, *const tk = token(s, &tklen, &ntok);

    if (!tk || ntok)
    {
        kprintf("Syntax: p:<index>\n");
        return 1;
    }
    else if (parse_num(tk, &index))
    {
        kprintf("Invalid value: %s\n", tk);
        return 1;
    }

    nw_dbg_param(&d->dbg, index);
    *next = (const struct nw_next){.fn = write_value, .user = p};
    return 0;
}

static int parse_type(const char *const s, enum nw_type *const out)
{
    switch (*s)
    {
        case 'i':
            *out = NW_TYPE_I32;
            break;

        case 'I':
            *out = NW_TYPE_I64;
            break;

        case 'f':
            *out = NW_TYPE_F32;
            break;

        case 'F':
            *out = NW_TYPE_F64;
            break;

        default:
            return -1;
    }

    return 0;
}

static int cmd_m(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    size_t tklen;
    struct bin_dbg *const d = &p->dbg;
    static const char exc[] = "Syntax: m:<i|I|f|F>:<addr>\n";
    const char *ntok, *tk = token(s, &tklen, &ntok);
    enum nw_type type;
    unsigned long addr;

    if (!tk || tklen != 1 || !ntok)
    {
        kprintf(exc);
        return 1;
    }
    else if (parse_type(tk, &type))
    {
        kprintf("Invalid type specifier: %.*s\n", (int)tklen, tk);
        return 1;
    }
    else if (!token(tk = ntok, &tklen, &ntok) || ntok)
    {
        kprintf(exc);
        return 1;
    }
    else if (parse_num(tk, &addr))
    {
        kprintf("Invalid address: %s\n", tk);
        return 1;
    }

    *next = (const struct nw_next){.fn = write_value, .user = p};
    nw_dbg_mem_load(&d->dbg, type, addr);
    return 0;
}

static int cmd_s(struct bin_proc *const p, const char *const s,
    struct nw_next *const next)
{
    struct bin_dbg *const d = &p->dbg;

    if (s)
    {
        kprintf("Syntax: s\n");
        return 1;
    }

    d->running = d->step = true;
    *next = (const struct nw_next){.fn = ok, .user = p};
    return 0;
}

static int parse_cmd(struct bin_proc *const p, const char *const tk,
    const size_t tklen, const char *const ntk, struct nw_next *const next)
{
    static const struct cmd
    {
        const char *cmd;
        int (*f)(struct bin_proc *, const char *, struct nw_next *);
    } cmds[] =
    {
        {.cmd = "b", .f = cmd_b},
        {.cmd = "bt", .f = cmd_bt},
        {.cmd = "c", .f = cmd_c},
        {.cmd = "d", .f = cmd_d},
        {.cmd = "g", .f = cmd_g},
        {.cmd = "l", .f = cmd_l},
        {.cmd = "m", .f = cmd_m},
        {.cmd = "p", .f = cmd_p},
        {.cmd = "s", .f = cmd_s}
    };

    for (size_t i = 0; i < sizeof cmds / sizeof *cmds; i++)
    {
        const struct cmd *const c = &cmds[i];

        if (!strncmp(c->cmd, tk, tklen))
            return c->f(p, ntk, next);
    }

    kprintf("Unknown command %.*s\n", (int)tklen, tk);
    return 1;
}

static int parse(struct bin_proc *const p, struct nw_next *const next)
{
    struct bin_dbg *const dbg = &p->dbg;
    struct dynstr *const d = &dbg->dstr;
    const char *const s = d->str, *tk, *ntk;
    size_t tklen;

    if (!s || !(tk = token(s, &tklen, &ntk)) || !tklen)
        return 1;

    return parse_cmd(p, tk, tklen, ntk, next);
}

static enum nw_state check_lf(void *const user, struct nw_next *const next)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const dbg = &p->dbg;
    struct dynstr *const d = &dbg->dstr;
    const char b = dbg->b;
    const size_t max = 128;

    if (b == ';')
    {
        dynstr_free(d);

        if (read_string(p, next))
            return NW_FATAL;
    }
    else if (b == '\n')
    {
        const int n = parse(p, next);

        if (n < 0)
            goto failure;
        else if (n)
        {
            dynstr_free(d);
            *next = (const struct nw_next){.fn = nok, .user = p};
        }
    }
    else if (d->len < max)
        if (dynstr_append(d, "%c", b) || read_byte(p))
            goto failure;

    return NW_AGAIN;

failure:
    dynstr_free(d);
    return NW_FATAL;
}

static int read_string(struct bin_proc *const p, struct nw_next *const next)
{
    struct bin_dbg *const d = &p->dbg;

    dynstr_init(&d->dstr);

    if (read_byte(p))
        return -1;

    *next = (const struct nw_next){.fn = check_lf, .user = p};
    return 0;
}

static enum nw_state check_start(void *const user, struct nw_next *const next)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const d = &p->dbg;

    if (d->b == ';')
    {
        if (read_string(p, next))
            return NW_FATAL;
    }
    else if (read_byte(p))
        return NW_FATAL;

    return NW_AGAIN;
}

enum nw_state bin_proc_pc(const long offset, struct nw_next *const next,
    void *const user)
{
    struct bin_proc *const p = user;
    struct bin_dbg *const d = &p->dbg;

    if (!d->running || d->step)
    {
        d->step = d->running = false;
        d->offset = offset;
        *next = (const struct nw_next){.fn = trap, .user = p};
        return NW_AGAIN;
    }
    else
    {
        for (size_t i = 0; i < d->n_bkpt; i++)
            if (offset == d->bkpt[i])
            {
                kprintf("Reached breakpoint %zu at offset %#lx\n", i, offset);
                d->running = d->step = false;
                d->offset = offset;
                *next = (const struct nw_next){.fn = trap, .user = p};
                return NW_AGAIN;
            }

        const struct fs_read r =
        {
            .fd = &d->fd,
            .buf = &d->b,
            .n = sizeof d->b
        };

        const int n = aio_read_nb(&r);

        if (n < 0)
        {
            kprintf("%s: aio_read_nb failed\n", __func__);
            return NW_FATAL;
        }
        else if (n && d->b == -1)
            d->running = false;
    }

    return NW_OK;
}
