/// sforth, an implementation of forth
/// Copyright (C) 2024  Clyne Sullivan <clyne@bitgloo.com>
///
/// This program is free software: you can redistribute it and/or modify it
/// under the terms of the GNU 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 General Public License for
/// more details.
///
/// You should have received a copy of the GNU General Public License along
/// with this program.  If not, see <http://www.gnu.org/licenses/>.

#ifndef SFORTH_HPP
#define SFORTH_HPP

#include <algorithm>
#include <array>
#include <bit>
#include <charconv>
#include <cstdint>
#include <cstddef>
#include <iterator>
#include <span>
#include <string_view>
#include <tuple>
#include <utility>

using cell = std::intptr_t;
using addr = std::uintptr_t;
using func = void (*)(const void *);

struct word_base;

struct word_list
{
    const word_base *next;

    constexpr word_list(const word_base *n = nullptr): next{n} {}

    std::optional<const word_base *> get(std::string_view sv) const;

    static constexpr auto parse(const char *source, std::size_t& sourcei) -> std::string_view {
        const std::string_view sv {source};

        const auto e = sv.find_first_of(" \t\r\n", sourcei);
        const auto word = e != std::string_view::npos ?
            sv.substr(sourcei, e - sourcei) : sv.substr(sourcei);
    
        sourcei = sv.find_first_not_of(" \t\r\n", e);
        return word;
    }
};

struct word_base : public word_list
{
    static constexpr addr immediate = 1 << 8;

    addr flags_len;

    constexpr word_base(const word_base *n, addr fl):
        word_list{n}, flags_len{fl} {}

    std::string_view name() const {
        return {std::bit_cast<const char *>(this + 1)};
    }

    const func *body() const {
        const auto ptr = std::bit_cast<const std::uint8_t *>(this + 1);
        const auto fptr = ptr + (flags_len & 0xFF);
        return std::bit_cast<const func *>(fptr);
    }

    constexpr void make_immediate() {
        flags_len |= immediate;
    }
};

template<unsigned N>
struct ctstring {
    char data[N];

    consteval ctstring(const char (&s)[N]) {
        std::copy(s, s + N, data);
    }
    consteval operator const char *() const {
        return data;
    }
    consteval auto size() const {
        return N;
    }
};

template<ctstring Name, ctstring Body, auto& Prev>
struct comp_word : public word_base
{
    constexpr static auto N = (sizeof(Name) + sizeof(cell) - 1) & ~(sizeof(cell) - 1);
    static constexpr std::size_t B = 
        [] {
            std::size_t b = 1;
            std::string_view sv {Body.data};
            auto sourcei = sv.find_first_not_of(" \t\r\n");
            while (sourcei != std::string_view::npos) {
                const auto word = parse(Body.data, sourcei);

                b++;
                if (!Prev.get_ct(word))
                    b++;
            }
            return b;
        }();

    union bodyt {
        const func *f;
        cell c;
    };

    std::array<char, N> namebuf;
    const func prologue;
    std::array<bodyt, B> bodybuf;

    consteval const func *get_ct(std::string_view name) const {
        if (name == std::string_view{Name.data})
            return &prologue;
        else
            return Prev.get_ct(name);
    }

    consteval comp_word(const func prol, addr flags = 0):
        word_base{&Prev, N | flags}, namebuf{}, prologue{prol}, bodybuf{}
    {
        std::copy(Name.data, Name.data + sizeof(Name), namebuf.data());

        auto bptr = bodybuf.begin();
        std::string_view sv {Body};
        auto sourcei = sv.find_first_not_of(" \t\r\n");
        while (sourcei != std::string_view::npos) {
            const auto word = parse(Body, sourcei);

            auto w = get_ct(word);
            if (w) {
                bptr->f = get_ct(word);
                bptr++;
            } else {
                cell n;
                std::from_chars(word.cbegin(), word.cend(), n, 10);

                bptr->f = get_ct("_lit");
                bptr++;
                bptr->c = n;
                bptr++;
            }
        }
    }
};

template<ctstring Name, auto *Prev = (const word_base *)nullptr>
struct native_word : public word_base
{
    constexpr static auto N = (sizeof(Name) + sizeof(cell) - 1) & ~(sizeof(cell) - 1);
    std::array<char, N> namebuf;
    func body;

    consteval const func *get_ct(std::string_view name) const {
        if (name == std::string_view{Name.data})
            return &body;
        else if constexpr (Prev != nullptr)
            return Prev->get_ct(name);
        else
            return nullptr;
    }

    consteval native_word(func bod, addr flags = 0):
        word_base{Prev, N | flags}, namebuf{}, body{bod}
    {
        std::copy(Name.data, Name.data + sizeof(Name), namebuf.data());
    }
};

struct forth : public word_list
{
    static constexpr bool enable_exceptions = true;
    static constexpr int data_size = 16;
    static constexpr int return_size = 16;

    static constexpr auto npos = std::string_view::npos;

    enum class error {
        init_error,
        parse_error,
        execute_error,
        dictionary_overflow,
        word_not_found,
        stack_underflow,
        stack_overflow,
        return_stack_underflow,
        return_stack_overflow,
        compile_only_word
    };

    template<error Err>
    static inline void assert(bool condition) {
        if constexpr (enable_exceptions) {
            if (!condition)
                throw Err;
        }
    }

    void push(cell v) {
        assert<error::stack_overflow>(sp != dstack.begin());
        *--sp = v;
    }

    void push(cell v, auto... vs) {
        push(v); (push(vs), ...);
    }

    void rpush(func *v) {
        assert<error::return_stack_overflow>(rp != rstack.begin());
        *--rp = v;
    }

    cell& top() {
        assert<error::stack_underflow>(sp != dstack.end());
        return *sp;
    }

    cell pop() {
        assert<error::stack_underflow>(sp != dstack.end());
        return *sp++;
    }

    auto rpop() -> func * {
        assert<error::return_stack_underflow>(rp != rstack.end());
        return *rp++;
    }

    template<int N>
    auto pop() {
        static_assert(N > 0, "pop<N>() with N <= 0");

        auto t = std::tuple {pop()};
        if constexpr (N > 1)
            return std::tuple_cat(t, pop<N - 1>());
        else
            return t;
    }

    forth& add(std::string_view name, func entry = nullptr) {
        const auto namesz = (name.size() + 1 + sizeof(cell) - 1) & ~(sizeof(cell) - 1);
        const auto size = (sizeof(word_base) + namesz) / sizeof(cell);

        assert<error::parse_error>(!name.empty());
        //assert<error::dictionary_overflow>(state->here + size < &dictionary.back());

        const auto h = std::exchange(here, here + size);
        next = new (h) word_base (next, namesz);
        std::copy(name.begin(), name.end(),
            std::bit_cast<char *>(h) + sizeof(word_base));
        if (entry)
            *here++ = std::bit_cast<cell>(entry);
        return *this;
    }

    auto parse() -> std::string_view {
        return word_list::parse(source, sourcei);
    }

    void parse_line(std::string_view sv) {
        source = sv.data();
        sourcei = sv.find_first_not_of(" \t\r\n");

        while (sourcei != npos) {
            const auto word = parse();

            if (auto ent = get(word); !ent) {
                cell n;
                const auto [p, e] = std::from_chars(word.cbegin(), word.cend(),
                    n, base);

                assert<error::word_not_found>(e == std::errc() && p == word.cend());

                push(n);

                if (compiling)
                    execute((*get("literal"))->body());
            } else {
                auto body = (*ent)->body();

                if (compiling && ((*ent)->flags_len & word_base::immediate) == 0) {
                    *here++ = std::bit_cast<cell>(body);
                } else {
                    execute(body);
                }
            }
        }
    }

    void execute(const func *body) {
        assert<error::execute_error>(body && *body);
        (*body)(body);
    }

    template<forth **fthp>
    static void prologue(const void *bodyf) {
        static auto& fth = **fthp;

        auto body = (func *)bodyf;
        fth.rpush(fth.ip);

        for (fth.ip = body + 1; *fth.ip; fth.ip++)
            fth.execute(std::bit_cast<func *>(*fth.ip));

        fth.ip = fth.rpop();
    }

    template<forth** fthp>
    static void initialize(cell *end_value)
    {
        assert<error::init_error>(*fthp);

        static auto& fth = **fthp;

        constexpr static func lit_impl = [](auto) {
            auto ptr = std::bit_cast<cell *>(++fth.ip);
            fth.push(*ptr);
        };
        auto f_dict   = [](auto) { fth.push(std::bit_cast<cell>(&fth)); };
        auto f_add    = [](auto) { fth.top() += fth.pop(); };
        auto f_minus  = [](auto) { fth.top() -= fth.pop(); };
        auto f_times  = [](auto) { fth.top() *= fth.pop(); };
        auto f_divide = [](auto) { fth.top() /= fth.pop(); };
        auto f_mod    = [](auto) { fth.top() %= fth.pop(); };
        auto f_bitand = [](auto) { fth.top() &= fth.pop(); };
        auto f_bitor  = [](auto) { fth.top() |= fth.pop(); };
        auto f_bitxor = [](auto) { fth.top() ^= fth.pop(); };
        auto f_lshift = [](auto) { fth.top() <<= fth.pop(); };
        auto f_rshift = [](auto) { fth.top() >>= fth.pop(); };
        auto f_lbrac  = [](auto) { fth.compiling = false; };
        auto f_rbrac  = [](auto) { fth.compiling = true; };
        auto f_imm    = [](auto) {
            const_cast<word_base *>(fth.next)->make_immediate(); };
        auto f_lit    = [](auto) {
            //assert<error::compile_only_word>(fth.compiling);
            *fth.here++ = std::bit_cast<cell>(&lit_impl);
            *fth.here++ = fth.pop(); };
        auto f_peek = [](auto) { fth.push(*std::bit_cast<cell *>(fth.pop())); };
        auto f_poke = [](auto) {
            auto [p, v] = fth.pop<2>();
            *std::bit_cast<cell *>(p) = v; };
        auto f_cpeek = [](auto) { fth.push(*std::bit_cast<char *>(fth.pop())); };
        auto f_cpoke = [](auto) {
            auto [p, v] = fth.pop<2>();
            *std::bit_cast<char *>(p) = v; };
        auto f_swap = [](auto) { auto [a, b] = fth.pop<2>(); fth.push(a, b); };
        auto f_drop = [](auto) { fth.pop(); };
        auto f_dup  = [](auto) { fth.push(fth.top()); };
        auto f_rot  = [](auto) { auto [a, b, c] = fth.pop<3>(); fth.push(b, a, c); };
        auto f_eq   = [](auto) { auto v = fth.pop(); fth.top() = -(fth.top() == v); };
        auto f_lt   = [](auto) { auto v = fth.pop(); fth.top() = -(fth.top() < v); };
        auto f_tick = [](auto) {
            auto w = fth.parse();

            if (auto g = fth.get(w); g)
                fth.push(std::bit_cast<cell>((*g)->body()));
            else
                fth.push(0); };
        auto f_colon = [](auto) {
            const auto prologue = forth::prologue<fthp>;
            auto w = fth.parse();
            fth.add(w);
            *fth.here++ = std::bit_cast<cell>(prologue);
            fth.compiling = true; };
        auto f_semic = [](auto) { *fth.here++ = 0; fth.compiling = false; };
        auto f_comm  = [](auto) { fth.sourcei = npos; };
        auto f_cell  = [](auto) { fth.push(sizeof(cell)); };
        auto f_jmp   = [](auto) {
            auto ptr = ++fth.ip;
            fth.ip = *std::bit_cast<func **>(ptr) - 1;
        };
        auto f_jmp0  = [](auto) {
            auto ptr = ++fth.ip;

            if (fth.pop() == 0)
                fth.ip = *std::bit_cast<func **>(ptr) - 1;
        };
        auto f_postpone = [](auto) {
            assert<error::compile_only_word>(fth.compiling);

            auto w = fth.parse();
            auto g = fth.get(w);

            assert<error::word_not_found>(g.has_value());

            *fth.here++ = std::bit_cast<cell>((*g)->body());
        };

        constexpr static native_word<"_d">             w_dict  {f_dict};
        constexpr static native_word<"_lit", &w_dict>  w_liti  {lit_impl};
        constexpr static native_word<"swap", &w_liti>  w_swap  {f_swap};
        constexpr static native_word<"drop", &w_swap>  w_drop  {f_drop};
        constexpr static native_word<"dup",  &w_drop>  w_dup   {f_dup};
        constexpr static native_word<"rot",  &w_dup>   w_rot   {f_rot};
        constexpr static native_word<"+",    &w_rot>   w_add   {f_add};
        constexpr static native_word<"-",    &w_add>   w_minus {f_minus};
        constexpr static native_word<"*",    &w_minus> w_times {f_times};
        constexpr static native_word<"/",    &w_times> w_divid {f_divide};
        constexpr static native_word<"mod",  &w_divid> w_mod   {f_mod};
        constexpr static native_word<"and",  &w_mod>   w_and   {f_bitand};
        constexpr static native_word<"or",   &w_and>   w_or    {f_bitor};
        constexpr static native_word<"xor",  &w_or>    w_xor   {f_bitxor}; 
        constexpr static native_word<"lshift", &w_xor> w_lsh   {f_lshift};
        constexpr static native_word<"rshift", &w_lsh> w_rsh   {f_rshift};
        constexpr static native_word<"[",    &w_rsh>   w_lbrac {f_lbrac,
            word_base::immediate};
        constexpr static native_word<"]",    &w_lbrac> w_rbrac {f_rbrac};
        constexpr static native_word<"immediate", &w_rbrac> w_imm {f_imm};
        constexpr static native_word<"literal", &w_imm> w_lit  {f_lit,
            word_base::immediate};
        constexpr static native_word<"@",    &w_lit>   w_peek  {f_peek};
        constexpr static native_word<"!",    &w_peek>  w_poke  {f_poke};
        constexpr static native_word<"c@",   &w_poke>  w_cpeek {f_cpeek};
        constexpr static native_word<"c!",   &w_cpeek> w_cpoke {f_cpoke};
        constexpr static native_word<"=",    &w_cpoke> w_eq    {f_eq};
        constexpr static native_word<"<",    &w_eq>    w_lt    {f_lt};
        constexpr static native_word<"\'",   &w_lt>    w_tick  {f_tick};
        constexpr static native_word<":",    &w_tick>  w_colon {f_colon};
        constexpr static native_word<";",    &w_colon> w_semic {f_semic,
            word_base::immediate};
        constexpr static native_word<"\\",   &w_semic> w_comm  {f_comm,
            word_base::immediate};
        constexpr static native_word<"cell", &w_comm>  w_cell  {f_cell};
        constexpr static native_word<"_jmp", &w_cell>  w_jmp   {f_jmp};
        constexpr static native_word<"_jmp0", &w_jmp>  w_jmp0  {f_jmp0};
        constexpr static native_word<"postpone", &w_jmp0> w_postp {f_postpone,
            word_base::immediate};
        constexpr static comp_word<"cell+", "cell +", w_postp> w_cellp
            {forth::prologue<fthp>};
        constexpr static comp_word<"cells", "cell *", w_cellp> w_cells
            {forth::prologue<fthp>};
        constexpr static comp_word<"char+", "1 +", w_cells> w_charp
            {forth::prologue<fthp>};
        constexpr static comp_word<"1+", "1 +", w_charp> w_inc
            {forth::prologue<fthp>};
        constexpr static comp_word<"1-", "1 -", w_inc> w_dec
            {forth::prologue<fthp>};

        fth.next = &w_dec;
        fth.end = end_value;
    }

    static auto error_string(error err) noexcept -> std::string_view {
        using enum error;
        switch (err) {
        case init_error:             return "init error";
        case parse_error:            return "parse error";
        case execute_error:          return "execute error";
        case dictionary_overflow:    return "dictionary overflow";
        case word_not_found:         return "word not found";
        case stack_underflow:        return "stack underflow";
        case stack_overflow:         return "stack overflow";
        case return_stack_underflow: return "return stack underflow";
        case return_stack_overflow:  return "return stack overflow";
        case compile_only_word:      return "compile only word";
        default:                     return "unknown error";
        }
    }

    constexpr forth() {
        sp = dstack.end();
        rp = rstack.end();
    }

    cell *sp;
    func **rp;
    func *ip = nullptr;
    cell *here = std::bit_cast<cell *>(this + 1);
    const char *source = nullptr;
    std::size_t sourcei = npos;
    cell compiling = false;
    cell *end = nullptr;
    cell base = 10;
    std::array<cell, data_size> dstack;
    std::array<func *, return_size> rstack;
};

std::optional<const word_base *> word_list::get(std::string_view sv) const
{
    for (auto lt = next; lt; lt = lt->next) {
        if (sv == lt->name())
            return lt;
    }

    return {};
}

//static_assert(offsetof(word_base, flags_len) == 1 * sizeof(cell));
//static_assert(offsetof(forth, sp)        == 1 * sizeof(cell));
//static_assert(offsetof(forth, rp)        == 2 * sizeof(cell));
//static_assert(offsetof(forth, ip)        == 3 * sizeof(cell));
//static_assert(offsetof(forth, here)      == 4 * sizeof(cell));
//static_assert(offsetof(forth, source)    == 5 * sizeof(cell));
//static_assert(offsetof(forth, sourcei)   == 6 * sizeof(cell));
//static_assert(offsetof(forth, compiling) == 7 * sizeof(cell));
//static_assert(offsetof(forth, end)       == 8 * sizeof(cell));
//static_assert(offsetof(forth, base)      == 9 * sizeof(cell));

#endif // SFORTH_HPP