#include "pr_wren.hpp"

#include <cstring>
#include <filesystem>
#include <fstream>
#include <iosfwd>
#include <sstream>
#include <sys/stat.h>

#include "quakedef.hpp"
#include "q_stdinc.hpp"
#include "console.hpp"
#include "sys.hpp"
#include "vec.hpp"
#include "../wren/vm/wren_debug.h"


namespace {
    WrenVM *_vm;
    int entity_counter{0};
    std::map<std::string, WrenHandle *> class_handles{};
    std::map<std::string, WrenHandle *> function_handles{};
    std::map<int, wren::entvars *> entity_vars{};
    bool silentErrors{false};

    void writeFn(WrenVM *vm, const char *text) {
        console::info("%s", text);
    }

    void errorFn(WrenVM *vm, WrenErrorType type, const char *module, int line,
                 const char *message) {
        Host_Error("%s:%d - %s", module, line, message);
    }

    void _set_size(WrenVM *vm) {
        wrenEnsureSlots(vm, 7);
        auto id = static_cast<int>(wrenGetSlotDouble(vm, 1));
        auto vars = entity_vars[id];
        wrenGetListElement(vm, 2, 0, 4);
        wrenGetListElement(vm, 2, 1, 5);
        wrenGetListElement(vm, 2, 2, 6);
        const auto mins = math::vec3f{
            static_cast<float>(wrenGetSlotDouble(vm, 4)), static_cast<float>(wrenGetSlotDouble(vm, 5)),
            static_cast<float>(wrenGetSlotDouble(vm, 6))
        };
        wrenGetListElement(vm, 3, 0, 4);
        wrenGetListElement(vm, 3, 1, 5);
        wrenGetListElement(vm, 3, 2, 6);
        const auto maxs = math::vec3f{
            static_cast<float>(wrenGetSlotDouble(vm, 4)), static_cast<float>(wrenGetSlotDouble(vm, 5)),
            static_cast<float>(wrenGetSlotDouble(vm, 6))
        };

        vars->mins = mins;
        vars->maxs = maxs;
    }

    // error(err)
    void _error(WrenVM *vm) {
        const std::string err = wrenGetSlotString(vm, 1);

        console::info("WREN ERROR: %s\n", err.c_str());
        // TODO: see if there's literally any way to get the call stack
        Host_Error("Program error");
    }

    // normalize(vec)
    void _normalize(WrenVM *vm) {
        wrenEnsureSlots(vm, 5);
        wrenGetListElement(vm, 1, 0, 2);
        wrenGetListElement(vm, 1, 1, 3);
        wrenGetListElement(vm, 1, 2, 4);
        auto vec = math::vec3f{
            static_cast<float>(wrenGetSlotDouble(vm, 2)), static_cast<float>(wrenGetSlotDouble(vm, 3)),
            static_cast<float>(wrenGetSlotDouble(vm, 4))
        };
        vec = vec.normalize();
        wrenEnsureSlots(vm, 4);
        wrenSetSlotNewList(vm, 0);
        wrenSetSlotDouble(vm, 1, vec.x);
        wrenSetSlotDouble(vm, 2, vec.y);
        wrenSetSlotDouble(vm, 3, vec.z);
        wrenInsertInList(vm, 0, 0, 1);
        wrenInsertInList(vm, 0, 1, 2);
        wrenInsertInList(vm, 0, 2, 3);
    }

    void _remove(WrenVM *vm) {
        const auto id = static_cast<int>(wrenGetSlotDouble(vm, 1));
        const auto handle = sv.wentities[id];
        sv.wentities[id] = nullptr;
        entity_vars[id]->free = true;
        if (handle != nullptr) {
            wrenReleaseHandle(vm, handle);
        }
    }

    void _precache_model(WrenVM *vm) {
        const auto model = wrenGetSlotString(vm, 1);

        for (auto i = 0; i < MAX_MODELS; i++) {
            if (sv.model_precache[i] == nullptr) {
                sv.model_precache[i] = model;
                sv.models[i] = Mod_ForName(model, true);
                return;
            }
            if (std::strcmp(sv.model_precache[i], model) == 0) {
                return;
            }
        }
    }

    // preache_sound(sound)
    void _precache_sound(WrenVM *vm) {
        const auto sample = wrenGetSlotString(vm, 1);

        if (sv.state != ss_loading) {
            console::info("Precache can only be done in spawn functions");
            Host_Error("Wren Error");
        }

        for (auto &i: sv.sound_precache) {
            if (i == nullptr) {
                i = sample;
                break;
            }
            if (std::strcmp(i, sample) == 0) {
                return;
            }
        }
    }

    void _find_radius(WrenVM *vm) {
        wrenEnsureSlots(vm, 5);
        auto radius = static_cast<float>(wrenGetSlotDouble(vm, 2));
        radius *= radius;
        wrenGetListElement(vm, 1, 0, 2);
        wrenGetListElement(vm, 1, 1, 3);
        wrenGetListElement(vm, 1, 2, 4);
        auto origin = math::vec3f{
            static_cast<float>(wrenGetSlotDouble(vm, 2)), static_cast<float>(wrenGetSlotDouble(vm, 3)),
            static_cast<float>(wrenGetSlotDouble(vm, 4))
        };

        wrenSetSlotNewList(vm, 0);
        auto count = 0;
        for (auto [id, ent]: sv.wentities) {
            float d;
            //
            // if (ent->v.solid == SOLID_NOT)
            //     continue;
            auto vars = entity_vars[id];

            d = origin.x - (vars->origin.x + (vars->mins.x + vars->maxs.x) * 0.5f);
            float lensq = d * d;
            if (lensq > radius)
                continue;
            d = origin.y - (vars->origin.y + (vars->mins.y + vars->maxs.y) * 0.5f);
            lensq += d * d;
            if (lensq > radius)
                continue;
            d = origin.z - (vars->origin.z + (vars->mins.z + vars->maxs.z) * 0.5f);
            lensq += d * d;
            if (lensq > radius)
                continue;

            wrenSetSlotHandle(vm, 1, ent);
            wrenInsertInList(vm, 0, count++, 1);
        }
    }

    // sound(entity, channel, sample, volume, attenuation)
    void _sound(WrenVM *vm) {
        const auto entity = static_cast<int>(wrenGetSlotDouble(vm, 1));
        const auto channel = static_cast<int>(wrenGetSlotDouble(vm, 2));
        std::string sample{};
        if (wrenGetSlotType(vm, 3) == WREN_TYPE_STRING) {
            sample = wrenGetSlotString(vm, 3);
        } else {
            return;
        }
        const auto volume = static_cast<int>(wrenGetSlotDouble(vm, 4) * 255);
        const auto attenuation = static_cast<float>(wrenGetSlotDouble(vm, 5));

        SV_StartSound(&sv.edicts[entity], channel, sample.c_str(), volume, attenuation);
    }

    net::msg *_write_destination(const int dest) {
        switch (dest) {
            case 0: // MSG_BROADCAST
                return &sv.datagram;

            case 1: {
                // MSG_ONE
                auto ent = PROG_TO_EDICT(pr_global_struct->msg_entity);
                auto entnum = NUM_FOR_EDICT(ent);
                if (entnum < 1 || entnum > svs.maxclients)
                    PR_RunError("WriteDest: not a client");
                return &svs.clients[entnum - 1].message;
            }

            case 2: // MSG_ALL
                return &sv.reliable_datagram;

            case 3: // MSG_INIT
                return &sv.signons[sv.current_signon];

            default:
                PR_RunError("_write_destination: bad destination");
                break;
        }
    }

    void _write_byte(WrenVM *vm) {
        auto dest = _write_destination(static_cast<int>(wrenGetSlotDouble(vm, 1)));
        const auto value = static_cast<byte>(wrenGetSlotDouble(vm, 2));

        dest->write_byte(value);
    }

    void _write_char(WrenVM *vm) {
        auto dest = _write_destination(static_cast<int>(wrenGetSlotDouble(vm, 1)));
        const auto value = static_cast<char>(wrenGetSlotDouble(vm, 2));

        dest->write_char(value);
    }

    void _write_coord(WrenVM *vm) {
        auto dest = _write_destination(static_cast<int>(wrenGetSlotDouble(vm, 1)));
        const auto value = static_cast<float>(wrenGetSlotDouble(vm, 2));

        dest->write_coord(value, static_cast<net::rmq_flags>(sv.protocolflags));
    }

    void _time(WrenVM *vm) {
        wrenSetSlotDouble(vm, 0, sv.time);
    }

    void _deathmatch(WrenVM *vm) {
        wrenSetSlotBool(vm, 0, deathmatch.value != 0.0);
    }

    void _entvars_allocate(WrenVM *vm) {
        auto vars = static_cast<wren::entvars *>(wrenSetSlotNewForeign(vm, 0, 0, sizeof(wren::entvars)));
        vars->id = entity_counter++;
        vars->origin = math::ORIGIN;
        vars->next_think = 0;
        entity_vars[vars->id] = vars;
    }

    void _entvars_id(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));

        wrenSetSlotDouble(vm, 0, vars->id);
    }

    void _entvars_origin(WrenVM *vm) {
        wrenEnsureSlots(vm, 4);

        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));

        wrenSetSlotNewList(vm, 0);
        wrenSetSlotDouble(vm, 1, vars->origin.x);
        wrenSetSlotDouble(vm, 2, vars->origin.y);
        wrenSetSlotDouble(vm, 3, vars->origin.z);
        wrenInsertInList(vm, 0, 0, 1);
        wrenInsertInList(vm, 0, 1, 2);
        wrenInsertInList(vm, 0, 2, 3);
    }

    void _entvars_set_origin(WrenVM *vm) {
        wrenEnsureSlots(vm, 5);

        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));

        wrenGetListElement(vm, 1, 0, 2);
        wrenGetListElement(vm, 1, 1, 3);
        wrenGetListElement(vm, 1, 2, 4);
        const auto origin = math::vec3f{
            static_cast<float>(wrenGetSlotDouble(vm, 2)), static_cast<float>(wrenGetSlotDouble(vm, 3)),
            static_cast<float>(wrenGetSlotDouble(vm, 4))
        };

        vars->origin = origin;
    }

    void _entvars_next_think(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        wrenEnsureSlots(vm, 1);
        wrenSetSlotDouble(vm, 0, vars->next_think);
    }

    void _entvars_set_next_think(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        const auto next_think = static_cast<float>(wrenGetSlotDouble(vm, 1));
        vars->next_think = next_think;
    }

    void _entvars_target_name(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        if (!vars->target_name.empty()) {
            wrenSetSlotString(vm, 0, vars->target_name.c_str());
        } else {
            wrenSetSlotNull(vm, 0);
        }
    }

    void _entvars_class_name(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        if (!vars->class_name.empty()) {
            wrenSetSlotString(vm, 0, vars->class_name.c_str());
        } else {
            wrenSetSlotNull(vm, 0);
        }
    }

    void _entvars_spawn_flags(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        wrenSetSlotDouble(vm, 0, vars->spawn_flags);
    }

    void _entvars_set_spawn_flags(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        const auto spawn_flags = static_cast<int>(wrenGetSlotDouble(vm, 1));
        vars->spawn_flags = spawn_flags;
    }

    void _entvars_sounds(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        wrenSetSlotDouble(vm, 0, vars->sounds);
    }

    void _entvars_noise(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        if (!vars->noise.empty()) {
            wrenSetSlotString(vm, 0, vars->noise.c_str());
        } else {
            wrenSetSlotNull(vm, 0);
        }
    }

    void _entvars_set_noise(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        vars->noise = wrenGetSlotString(vm, 1);
    }

    std::map<std::string, wren::typing> _entvars_type_map{
        {"targetname", wren::typing::STRING},
        {"classname", wren::typing::STRING},
        {"origin", wren::typing::VECTOR},
        {"spawnflags", wren::typing::INT},
        {"sounds", wren::typing::FLOAT}
    };


    void _entvars_field_setter(WrenVM *vm) {
        wrenEnsureSlots(vm, 9);
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        std::string name = wrenGetSlotString(vm, 1);
        if (name == "classname") {
            vars->class_name = wrenGetSlotString(vm, 2);
            return;
        }
        if (name == "targetname") {
            vars->target_name = wrenGetSlotString(vm, 2);
            return;
        }
        if (name == "origin") {
            vars->origin = math::vec3f{};
            wrenGetListElement(vm, 2, 0, 3);
            wrenGetListElement(vm, 2, 1, 4);
            wrenGetListElement(vm, 2, 2, 5);
            vars->origin.x = static_cast<float>(wrenGetSlotDouble(vm, 3));
            vars->origin.y = static_cast<float>(wrenGetSlotDouble(vm, 4));
            vars->origin.z = static_cast<float>(wrenGetSlotDouble(vm, 5));
            return;
        }
        if (name == "spawnflags") {
            vars->spawn_flags = static_cast<int>(wrenGetSlotDouble(vm, 2));
            return;
        }
        if (name == "sounds") {
            vars->sounds = static_cast<float>(wrenGetSlotDouble(vm, 2));
            return;
        }
        Host_Error("_entvars_field_setter: tried to set %s, but no setter defined", name.c_str());
    }

    void _entvars_field_type(WrenVM *vm) {
        const auto vars = static_cast<wren::entvars *>(wrenGetSlotForeign(vm, 0));
        const std::string name = wrenGetSlotString(vm, 1);
        auto kind = wren::typing::UNDEFINED;
        if (_entvars_type_map.contains(name)) {
            kind = _entvars_type_map[name];
        }
        wrenSetSlotDouble(vm, 0, static_cast<int>(kind));
    }

    WrenForeignClassMethods bindForeignClassFn(
        WrenVM *vm, const char *module, const char *className) {
        const std::string modname{module};
        const std::string classname{className};
        WrenForeignClassMethods methods{.allocate = nullptr, .finalize = nullptr};

        if (classname == "EntityVariables") {
            methods.allocate = _entvars_allocate;
        }

        return methods;
    }


    std::map<std::string, std::map<std::string, WrenForeignMethodFn> > exported_functions{
        {
            "EntityVariables",
            {
                {"id", _entvars_id},
                {"origin", _entvars_origin},
                {"origin=(_)", _entvars_set_origin},
                {"next_think", _entvars_next_think},
                {"next_think=(_)", _entvars_set_next_think},
                {"target_name", _entvars_target_name},
                {"class_name", _entvars_class_name},
                {"field_setter(_,_)", _entvars_field_setter},
                {"field_type(_)", _entvars_field_type},
                {"spawn_flags", _entvars_spawn_flags},
                {"spawn_flags=(_)", _entvars_set_spawn_flags},
                {"sounds", _entvars_sounds},
                {"noise", _entvars_noise},
                {"noise=(_)", _entvars_set_noise},
            },
        },
        {
            "Quake",
            {
                {"set_size(_,_,_)", _set_size},
                {"error(_)", _error},
                {"normalize(_)", _normalize},
                {"remove(_)", _remove},
                {"precache_model(_)", _precache_model},
                {"precache_sound(_)", _precache_sound},
                {"find_radius(_,_)", _find_radius},
                {"sound(_,_,_,_,_)", _sound},
                {"write_byte(_,_)", _write_byte},
                {"write_char(_,_)", _write_char},
                {"write_coord(_,_)", _write_coord},
                {"time", _time},
                {"deathmatch", _deathmatch},
            }
        },
    };

    WrenForeignMethodFn bindForeignMethodFn(WrenVM *vm, const char *module, const char *className, bool isStatic,
                                            const char *signature) {
        const std::string modname{module};
        const std::string classname{className};
        const std::string sig{signature};

        if (modname == "quake") {
            if (exported_functions.contains(classname)) {
                if (auto &clazz = exported_functions[classname]; clazz.contains(sig)) {
                    return clazz[sig];
                }
            }
        }
        return nullptr;
    }
}

void wren::init() {
    WrenConfiguration config;
    wrenInitConfiguration(&config);
    config.writeFn = writeFn;
    config.errorFn = errorFn;
    config.bindForeignMethodFn = bindForeignMethodFn;
    config.bindForeignClassFn = bindForeignClassFn;
    _vm = wrenNewVM(&config);
    function_handles["vars"] = nullptr;
    function_handles["new()"] = nullptr;
    function_handles["spawn()"] = nullptr;
    function_handles["think()"] = nullptr;
    function_handles["field_type(_)"] = nullptr;
    function_handles["field_setter(_,_)"] = nullptr;
}

std::optional<std::pair<int, WrenHandle *> > wren::create_entity(const std::string &classname,
                                                                 std::map<std::string, std::string> pairs) {
    if (!wrenHasVariable(_vm, "quake", classname.c_str())) {
        console::safe_print("No Wren spawn function for %s\n", classname.c_str());
        return std::nullopt;
    }

    wrenEnsureSlots(_vm, 2);
    if (!class_handles.contains(classname)) {
        wrenEnsureSlots(_vm, 1);
        wrenGetVariable(_vm, "quake", classname.c_str(), 0);
        class_handles[classname.c_str()] = wrenGetSlotHandle(_vm, 0);
    }

    wrenSetSlotHandle(_vm, 0, class_handles[classname]);
    wrenCall(_vm, function_handles["new()"]);
    const auto handle = wrenGetSlotHandle(_vm, 0);

    for (const auto &[key, value]: pairs) {
        set_field(handle, key, value);
    }

    entity_vars[entity_counter - 1]->next_think = sv.time;

    wrenSetSlotHandle(_vm, 0, handle);
    wrenCall(_vm, function_handles["spawn()"]);
    if (entity_vars[entity_counter - 1]->free) {
        return std::nullopt;
    }

    console::info("loaded %s\n", classname.c_str());

    return std::pair{entity_counter - 1, handle};
}


void wren::load_progs() {
    wrenNukeModules(_vm);
    std::filesystem::directory_iterator dir{"id1/wprog"};
    for (auto f: dir) {
        if (f.path().extension() == ".wren") {
            std::ifstream ifs{f.path()};
            std::stringstream ss{};
            ss << ifs.rdbuf();
            wrenInterpret(_vm, "quake", ss.str().c_str());
        }
    }

    for (const auto& [name, handle] : class_handles) {
        if (handle) wrenReleaseHandle(_vm, handle);
        wrenEnsureSlots(_vm, 1);
        wrenGetVariable(_vm, "quake", name.c_str(), 0);
        class_handles[name] = wrenGetSlotHandle(_vm, 0);
    }

    for (const auto &[signature, handle]: function_handles) {
        if (handle) wrenReleaseHandle(_vm, handle);
        function_handles[signature] = wrenMakeCallHandle(_vm, signature.c_str());
    }

    entity_vars.clear();
    entity_counter = 0;

    wrenEnsureSlots(_vm, 1);
    wrenGetVariable(_vm, "quake", "EntityVariables", 0);
    class_handles["EntityVariables"] = wrenGetSlotHandle(_vm, 0);
    wrenEnsureSlots(_vm, 1);
    wrenGetVariable(_vm, "quake", "Typing", 0);
    class_handles["Typing"] = wrenGetSlotHandle(_vm, 0);
}

void wren::run_think(const int index) {
    const auto ent = sv.wentities[index];
    if (ent == nullptr) {
        return;
    }

    auto vars = entity_vars[index];

    auto thinktime = vars->next_think;
    if (thinktime <= 0 || thinktime > sv.time + host_frametime)
        return;

    if (thinktime < sv.time)
        thinktime = sv.time;

    // ent->oldthinktime = thinktime;
    // ent->oldframe = ent->v.frame; //johnfitz

    vars->next_think = 0;

    wrenEnsureSlots(_vm, 1);
    wrenSetSlotHandle(_vm, 0, ent);
    wrenCall(_vm, function_handles["think()"]);
}

void wren::set_field(WrenHandle *ent, std::string key, std::string value) {
    wrenEnsureSlots(_vm, 9);
    wrenSetSlotHandle(_vm, 0, ent);
    wrenSetSlotString(_vm, 1, key.c_str());
    wrenCall(_vm, function_handles["field_type(_)"]);
    auto kind = static_cast<typing>(wrenGetSlotDouble(_vm, 0));
    switch (kind) {
        case typing::UNDEFINED:
            console::info("FIELD %s IS UNDEFINED\n", key.c_str());
            break;
        case typing::STRING:
            wrenSetSlotHandle(_vm, 0, ent);
            wrenSetSlotString(_vm, 1, key.c_str());
            wrenSetSlotString(_vm, 2, value.c_str());
            wrenCall(_vm, function_handles["field_setter(_,_)"]);
            break;
        case typing::INT: {
            wrenSetSlotHandle(_vm, 0, ent);
            wrenSetSlotString(_vm, 1, key.c_str());
            auto i = std::atoi(value.c_str());
            wrenSetSlotDouble(_vm, 2, i);
            wrenCall(_vm, function_handles["field_setter(_,_)"]);
            break;
        }
        case typing::FLOAT: {
            wrenSetSlotHandle(_vm, 0, ent);
            wrenSetSlotString(_vm, 1, key.c_str());
            auto i = static_cast<float>(std::atof(value.c_str()));
            wrenSetSlotDouble(_vm, 2, i);
            wrenCall(_vm, function_handles["field_setter(_,_)"]);
            break;
        }
        case typing::VECTOR: {
            math::vec3f res{};
            std::sscanf(value.c_str(), "%f %f %f", &res.x, &res.y, &res.z);
            wrenSetSlotHandle(_vm, 0, ent);
            wrenSetSlotString(_vm, 1, key.c_str());
            wrenSetSlotNewList(_vm, 2);
            wrenSetSlotDouble(_vm, 3, res.x);
            wrenSetSlotDouble(_vm, 4, res.y);
            wrenSetSlotDouble(_vm, 5, res.z);
            wrenInsertInList(_vm, 2, 0, 3);
            wrenInsertInList(_vm, 2, 1, 4);
            wrenInsertInList(_vm, 2, 2, 5);
            wrenCall(_vm, function_handles["field_setter(_,_)"]);
            break;
        }
    }
}

std::map<std::string, std::string> wren::parse_fields(std::istringstream &ss) {
    std::map<std::string, std::string> vars{};
    bool anglehack{false};


    while (true) {
        auto token = common::parse_token(ss);
        if (token.has_value() && token->front() == '}') {
            break;
        }
        if (ss.eof()) {
            Host_Error("wren::parse_fields: EOF without closing brace");
        }

        // anglehack is to allow QuakeEd to write single scalar angles
        // and allow them to be turned into vectors. (FIXME...)
        if (token == "angle") {
            token = "angles";
            anglehack = true;
        } else {
            anglehack = false;
        }

        // FIXME: change light to _light to get rid of this hack
        if (token == "light") {
            token = "light_maplev";
        }

        auto keyname = *token;

        // another hack to fix keynames with trailing spaces
        while (!keyname.empty() && keyname.back() == ' ') {
            keyname.pop_back();
        }

        // parse value
        // HACK: we allow truncation when reading the wad field,
        // otherwise maps using lots of wads with absolute paths
        // could cause a parse error
        token = common::parse_token(ss);
        if (ss.eof()) {
            Host_Error("wren::parse_fields: EOF without closing brace");
        }

        if (!token->empty() && token->front() == '}') {
            Host_Error("wren::parse_fields: closing brace without data");
        }


        // keynames with a leading underscore are used for utility comments,
        // and are immediately discarded by quake
        if (keyname[0] == '_') {
            continue;
        }

        if (anglehack) {
            token = std::format("0 {} 0", *token);
        }

        vars[keyname] = *token;
    }

    return vars;
}