#ifndef SCRIPT_ENGINE_HPP_
#define SCRIPT_ENGINE_HPP_

#include <Log/Log.hpp>
#include <Utilities/Extern.hpp>
#include <optional>
#include <string>
#include <unordered_map>
#include <VoxelSprite/VoxelSprite.hpp>

/**
 * @brief The execution environment for Lua scripts.
 */
class ScriptEngine
{
public:
    /**
     * @brief Special return type for async methods; either contains the result or 'yields'.
     */
    template<typename T>
    class Async final
    {
    private:
        std::optional<T> result;

        Async() {};
    
    public:
        Async(T &&t) : result(t) {}

        static Async<T> yield()
        {
            return Async();
        }

        bool hasYielded() const
        {
            return !result.has_value();
        }

        const T& get() const
        {
            if (!result.has_value())
            {
                crash("Attempted to get() from a yielded Async<T>.");
            }

            return result.value();
        }
    };

    template<>
    class Async<void> final
    {
    private:
        bool yielded;
    
    public:
        Async(bool yielded) : yielded(yielded) {}

        static Async<void> yield()
        {
            return Async(true);
        }

        static Async<void> done()
        {
            return Async(false);
        }

        bool hasYielded() const
        {
            return yielded;
        }
    };

private:
    lua_State *L;
    lua_State *co;
    int coref;

    template<typename TReturn>
    struct ReturnCodec
    {
        static int push(ScriptEngine &engine, const TReturn &value)
        {
            engine.push(value);
            return 1;
        }
    };

    template<>
    struct ReturnCodec<glm::ivec3>
    {
        static int push(ScriptEngine &engine, const glm::ivec3 &value)
        {
            engine.push(value.x);
            engine.push(value.y);
            engine.push(value.z);
            return 3;
        }
    };

    template<typename T>
    struct AsyncReturnCodec
    {
        static int push(ScriptEngine &engine, const Async<T> &value)
        {
            if (value.hasYielded())
            {
                return -1;
            }

            return ReturnCodec<T>::push(engine, value.get());
        }
    };

    template<>
    struct AsyncReturnCodec<void>
    {
        static int push(ScriptEngine &engine, const Async<void> &value)
        {
            if (value.hasYielded())
            {
                return -1;
            }

            return 0;
        }
    };

    template<typename T>
    struct ReturnCodec<Async<T>> : public AsyncReturnCodec<T> {};

    class AbstractCallWrapper
    {
    public:
        struct SavedCallFrame
        {
            virtual ~SavedCallFrame() = default;
            virtual int resume() = 0;
        };

        virtual int implementCall(SavedCallFrame* &context) = 0;
        virtual ~AbstractCallWrapper() = default;
    };

    template<typename TReturn, typename... TArgs>
    class ConcreteCallWrapper : public AbstractCallWrapper
    {
    private:
        ScriptEngine &engine;
        std::string funcName;
        std::function<TReturn(TArgs...)> callback;

        template<typename TNextArg, typename... TOtherArgs>
        bool readNextArgs(int nextArgIndex, TNextArg &nextArg, TOtherArgs&... otherArgs)
        {
            return engine.read(nextArg, nextArgIndex) &&
                readNextArgs(nextArgIndex + 1, otherArgs...);
        }

        bool readNextArgs(int nextArgIndex)
        {
            return true;
        }

        int doCall(TArgs... args)
        {
            if constexpr (std::is_same_v<TReturn, void>)
            {
                return doCallImplVoid(args...);
            }
            else
            {
                return doCallImplNonVoid(args...);
            }
        }

        int doCallImplNonVoid(TArgs... args)
        {
            TReturn result = callback(args...);
            return ReturnCodec<TReturn>::push(engine, result);
        }

        int doCallImplVoid(TArgs... args)
        {
            callback(args...);
            return 0;
        }

    public:
        struct SavedCallFrame final : public AbstractCallWrapper::SavedCallFrame
        {
            ConcreteCallWrapper<TReturn, TArgs...> &wrapper;
            std::tuple<TArgs...> args;

            SavedCallFrame(ConcreteCallWrapper<TReturn, TArgs...> &wrapper) :
                wrapper(wrapper)
            {}

            virtual int resume() override
            {
                return wrapper.cont(this);
            }
        };

        ConcreteCallWrapper(ScriptEngine &engine, std::string funcName, std::function<TReturn(TArgs...)> callback) :
            engine(engine), funcName(funcName), callback(callback)
        {}

        int cont(SavedCallFrame *frame)
        {
            return std::apply([this](TArgs&... args) { return doCall(args...); }, frame->args);
        }

        virtual int implementCall(AbstractCallWrapper::SavedCallFrame* &context) override
        {
            size_t argsExpected = std::tuple_size<std::tuple<TArgs...>>::value;
            size_t argsReceived = lua_gettop(engine.co);

            if (argsReceived != argsExpected)
            {
                Log::error(
                    formatString(
                        "ScriptEngine: '$+' expects exactly $+ arguments, $+ received",
                        funcName, argsExpected, argsReceived
                    )
                );
            }

            SavedCallFrame *frame = new SavedCallFrame(*this);
            if (!std::apply([this](TArgs&... args) { return readNextArgs(1, args...); }, frame->args))
            {
                Log::error(formatString("ScriptEngine: Invalid argument types passed to '$+'.", funcName));
                engine.push(nullptr);
                delete frame;
                return 1;
            }

            context = frame;
            return cont(frame);
        }
    };

    std::unordered_map<std::string, std::unique_ptr<AbstractCallWrapper>> funcs;

    static int luaCallback(lua_State *L);
    static int luaContCallback(lua_State *L, int status, lua_KContext ctx);

public:
    ScriptEngine();
    ~ScriptEngine();

    ScriptEngine(const ScriptEngine &) = delete;
    ScriptEngine(const ScriptEngine &&) = delete;

    void push(std::string str);
    void push(std::nullptr_t);
    void push(int value);
    void push(VoxelSprite::Color color);

    bool read(std::string &str, int index);
    bool read(int &value, int index);
    bool read(VoxelSprite::Color &color, int index);

    bool execFromString(std::string luaCode);
    bool execFromFile(std::string luaScriptPath);

    bool loadFromString(std::string luaCode);
    bool loadFromFile(std::string luaScriptPath);

    Async<bool> resume();
    bool exec();

    void setGlobal(std::string name);

    template<typename TReturn, typename... TArgs>
    void addFunc(std::string name, std::function<TReturn(TArgs...)> callback)
    {
        auto wrapper = std::make_unique<ConcreteCallWrapper<TReturn, TArgs...>>(*this, name, callback);
        lua_pushlightuserdata(co, reinterpret_cast<void*>(wrapper.get()));

        funcs[name] = std::move(wrapper);

        lua_pushcclosure(co, luaCallback, 1);
        lua_setglobal(co, name.c_str());
    }

    template<typename T>
    void addFunc(std::string name, T t)
    {
        addFunc(name, std::function(t));
    }

    template<typename T>
    T pop()
    {
        T result;
        read(result, -1);
        lua_pop(co, 1);
        return result;
    }

    template<typename TKey, typename TValue>
    void push(std::unordered_map<TKey, TValue> map)
    {
        lua_newtable(co);

        for (const auto &[key, value] : map)
        {
            push(key);
            push(value);
            lua_settable(co, -3);
        }
    }

    template<typename TKey, typename TValue>
    bool read(std::unordered_map<TKey, TValue> &map, int index)
    {
        bool ok = true;

        lua_pushvalue(co, index);
        lua_pushnil(co);

        while (lua_next(co, -2))
        {
            lua_pushvalue(co, -2);
            TKey key;
            read(key, -1);
            ok = read(map[key], -2) && ok;
            lua_pop(co, 2);
        }

        lua_pop(co, 1);
        return ok;
    }

    template<typename T>
    bool read(std::vector<T> &vec, int index)
    {
        bool ok = true;

        lua_pushvalue(co, index);

        vec.clear();

        int len = lua_rawlen(co, -1);
        vec.reserve(len);

        for (int i = 1; i <= len; ++i) {
            lua_rawgeti(co, -1, i);
            T value;
            if (!read(value, -1)) {
                ok = false;
            }
            vec.push_back(std::move(value));
            lua_pop(co, 1);
        }

        lua_pop(co, 1);
        return ok;
    }
};

#endif