PkmnLib/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp

#include "AngelScriptResolver.hpp"
#include <CreatureLib/Battling/Models/Creature.hpp>
#include <regex>
#include "../../../extern/angelscript_addons/scriptdictionary/scriptdictionary.h"
#include "../../../extern/angelscript_addons/scripthandle/scripthandle.h"
#include "../../../extern/angelscript_addons/scripthelper/scripthelper.h"
#include "../../../extern/angelscript_addons/scriptstdstring/scriptstdstring.h"
#include "../../Battling/PkmnScriptCategory.hpp"
#include "../../Battling/Pokemon/Pokemon.hpp"
#include "ByteCodeHandling/FileByteCodeStream.hpp"
#include "ByteCodeHandling/MemoryByteCodeStream.hpp"
#include "TypeRegistry/BasicScriptClass.hpp"
#include "TypeRegistry/Battling/RegisterBattleClass.hpp"
#include "TypeRegistry/Battling/RegisterBattleLibrary.hpp"
#include "TypeRegistry/Battling/RegisterExecutingAttack.hpp"
#include "TypeRegistry/Battling/RegisterParty.hpp"
#include "TypeRegistry/Battling/RegisterPokemonClass.hpp"
#include "TypeRegistry/Battling/RegisterTurnChoices.hpp"
#include "TypeRegistry/ConstString.hpp"
#include "TypeRegistry/HelperFile.hpp"
#include "TypeRegistry/Library/RegisterEffectParameter.hpp"
#include "TypeRegistry/Library/RegisterGrowthRateTypes.hpp"
#include "TypeRegistry/Library/RegisterItemTypes.hpp"
#include "TypeRegistry/Library/RegisterMoveTypes.hpp"
#include "TypeRegistry/Library/RegisterSpeciesTypes.hpp"
#include "TypeRegistry/Library/RegisterStaticLibraryTypes.hpp"
#include "TypeRegistry/Library/RegisterTypeLibrary.hpp"

CreatureLib::Battling::ScriptResolver* PkmnLib::Battling::BattleLibrary::CreateScriptResolver() {
    return new AngelScriptResolver();
}

static void TranslateException(asIScriptContext* ctx, void* /*userParam*/) {
    try {
        // Retrow the original exception so we can catch it again
        throw;
    } catch (ArbUt::Exception& e) {
        // Tell the VM the type of exception that occurred
        ctx->SetException(e.what());
    } catch (std::exception& e) {
        // Tell the VM the type of exception that occurred
        ctx->SetException(e.what());
    } catch (...) {
        // The callback must not allow any exception to be thrown, but it is not necessary
        // to explicitly set an exception string if the default exception string is sufficient
    }
}

void AngelScriptResolver::Initialize(CreatureLib::Battling::BattleLibrary* arg, bool includeStandard) {
    for (auto scriptCategory : ScriptCategoryHelper::GetValues()) {
        _typeDatabase.Insert(scriptCategory, {});
    }
    for (auto scriptCategory : PkmnScriptCategoryHelper::GetValues()) {
        _typeDatabase.Insert(static_cast<ScriptCategory>(scriptCategory), {});
    }

    auto library = (PkmnLib::Battling::BattleLibrary*)arg;
    _engine = asCreateScriptEngine();
    _engine->SetTranslateAppExceptionCallback(asFUNCTION(TranslateException), 0, asCALL_CDECL);

    int32_t r = _engine->SetMessageCallback(asFUNCTION(MessageCallback), nullptr, asCALL_CDECL);
    if (r < 0)
        throw ArbUt::Exception("Registering message callback failed.");

    _engine->SetEngineProperty(asEP_DISALLOW_EMPTY_LIST_ELEMENTS, true);
    _engine->SetEngineProperty(asEP_DISALLOW_VALUE_ASSIGN_FOR_REF_TYPE, false);
    _engine->SetEngineProperty(asEP_ALLOW_UNSAFE_REFERENCES, true);
    _engine->SetEngineProperty(asEP_ALWAYS_IMPL_DEFAULT_CONSTRUCT, true);
    _engine->SetEngineProperty(asEP_AUTO_GARBAGE_COLLECT, false);
    _engine->SetEngineProperty(asEP_REQUIRE_ENUM_SCOPE, true);
    _engine->SetEngineProperty(asEP_PROPERTY_ACCESSOR_MODE, 2);
    _engine->SetEngineProperty(asEP_COMPILER_WARNINGS, 2);

    if (includeStandard) {
        RegisterStdString(_engine);
        ConstStringRegister::Register(_engine);

        // Register Script Array type
        RegisterScriptArray(_engine, true);
        RegisterScriptHandle(_engine);
        RegisterScriptDictionary(_engine);
        r = _engine->RegisterGlobalFunction("void print(const string &in)", asFUNCTION(Print), asCALL_CDECL);
        if (r < 0)
            throw ArbUt::Exception("Registering print function failed.");
        r = _engine->RegisterGlobalFunction("void print(const constString &in)", asFUNCTION(PrintConst), asCALL_CDECL);
        if (r < 0)
            throw ArbUt::Exception("Registering print function failed.");
    }

    _builder.StartNewModule(_engine, "pkmn");
    _mainModule = _builder.GetModule();

    RegisterTypes();
    RegisterExceptionRoutines(_engine);

    if (library != nullptr) {
        auto& staticLib = library->GetStaticLib();
        _engine->RegisterGlobalProperty("const StaticLibrary@ StaticLib", (void*)staticLib.get());
    }

    _contextPool = new ContextPool(_engine);

    asPrepareMultithread();
}

void AngelScriptResolver::RegisterTypes() {
    // Register static library types
    RegisterSpeciesTypes::Register(_engine);
    RegisterItemTypes::Register(_engine);
    RegisterMoveTypes::Register(_engine);
    RegisterGrowthRateTypes::Register(_engine);
    RegisterTypeLibrary::Register(_engine);
    RegisterStaticLibraryTypes::Register(_engine);
    RegisterEffectParameter::Register(_engine);

    // Register battle types
    // Predeclare these two types, and declare their implementation later.
    [[maybe_unused]] int r = _engine->RegisterObjectType("Battle", 0, asOBJ_REF | asOBJ_NOCOUNT);
    Ensure(r >= 0);
    r = _engine->RegisterObjectType("BattleSide", 0, asOBJ_REF | asOBJ_NOCOUNT);
    Ensure(r >= 0);

    RegisterPokemonClass::Register(_engine);
    RegisterParty::Register(_engine);
    RegisterExecutingAttack::Register(_engine);
    RegisterTurnChoices::Register(_engine);
    RegisterBattleLibrary::Register(_engine);
    RegisterBattleClass::Register(_engine);

    // Register base script
    BasicScriptClass::Register(_engine);
}

void AngelScriptResolver::MessageCallback(const asSMessageInfo* msg, void*) {
    const char* type = "ERR ";
    if (msg->type == asMSGTYPE_WARNING)
        type = "WARN";
    else if (msg->type == asMSGTYPE_INFORMATION)
        type = "INFO";
    printf("%s (%d, %d) : %s : %s\n", msg->section, msg->row, msg->col, type, msg->message);
}

CreatureLib::Battling::BattleScript* AngelScriptResolver::LoadScript(ScriptCategory category,
                                                                     const ArbUt::StringView& scriptName) {
    ArbUt::Dictionary<ArbUt::StringView, AngelScriptTypeInfo*> innerDb;
    auto v = _typeDatabase.TryGet(category);
    if (!v.has_value()) {
        innerDb.Insert(scriptName, nullptr);
        _typeDatabase.Insert(category, innerDb);
        return nullptr;
    } else {
        innerDb = v.value();
    }
    auto t = innerDb.TryGet(scriptName);
    if (!t.has_value()) {
        innerDb.Insert(scriptName, nullptr);
        return nullptr;
    }
    if (t.value() == nullptr) {
        return nullptr;
    }
    auto ctx = _contextPool->RequestContext();
    auto obj = const_cast<AngelScriptTypeInfo*>(t.value().get())->Instantiate(ctx);
    _contextPool->ReturnContextToPool(ctx);
    return new AngelScriptScript(this, t.value(), obj, _contextPool);
}

CreatureLib::Battling::ItemUseScript* AngelScriptResolver::LoadItemScript(const CreatureLib::Library::Item* item) {
    auto v = this->_itemUseScripts.TryGet(item);
    if (v.has_value()) {
        return v.value();
    }
    if (!item->GetEffect().HasValue()) {
        return nullptr;
    }

    auto typeInfoOption = _itemUseTypes.TryGet(item->GetEffect().GetValue()->GetEffectName());
    if (!typeInfoOption.has_value()) {
        return nullptr;
    }

    auto* ctx = _contextPool->RequestContext();
    auto factory = typeInfoOption.value().get()->GetFactoryByIndex(0);
    ctx->Prepare(factory);
    auto result = ctx->Execute();
    if (result != asEXECUTION_FINISHED) {
        throw ArbUt::Exception("Instantiation failed.");
    }
    asIScriptObject* obj = *(asIScriptObject**)ctx->GetAddressOfReturnValue();
    obj->AddRef();
    auto scriptObject = new AngelScriptItemUseScript(obj, this);
    scriptObject->OnInitialize(item->GetEffect().GetValue()->GetParameters());
    _itemUseScripts.Insert(item, scriptObject);
    _contextPool->ReturnContextToPool(ctx);
    return scriptObject;
}

void AngelScriptResolver::FinalizeModule() {
    int r = _builder.BuildModule();
    if (r < 0)
        throw ArbUt::Exception("Building Script Module failed.");
    asUINT count = _mainModule->GetObjectTypeCount();
    std::regex metadataMatcher(R"(^\s*(\w+)([\w\s=]*)$)", std::regex_constants::icase);
    std::regex variableMatcher(R"(\s*(\w+)=(\w+))", std::regex_constants::icase);
    std::smatch base_match;

    auto pkmnScriptType = _mainModule->GetTypeInfoByName("PkmnScript");
    auto itemUseScriptType = _mainModule->GetTypeInfoByName("ItemUseScript");
    for (asUINT n = 0; n < count; n++) {
        auto typeInfo = _mainModule->GetObjectTypeByIndex(n);
        if (typeInfo->DerivesFrom(pkmnScriptType)) {
            auto metadata = _builder.GetMetadataForType(typeInfo->GetTypeId());
            for (size_t m = 0; m < metadata.size(); m++) {
                auto data = metadata[m];
                if (std::regex_match(data, base_match, metadataMatcher)) {
                    auto mt = base_match[1].str();
                    auto metadataKind = ArbUt::StringView(mt.c_str(), mt.length());
                    auto metadataVariables = base_match[2].str();
                    if (!std::regex_match(metadataVariables, base_match, variableMatcher)) {
                        continue;
                    }
                    ArbUt::StringView effectName;
                    for (size_t variableIndex = 1; variableIndex < base_match.size(); variableIndex += 2) {
                        if (ArbUt::StringView::CalculateHash(base_match[variableIndex].str().c_str()) == "effect"_cnc) {
                            auto val = base_match[variableIndex + 1].str();
                            effectName = ArbUt::StringView(val.c_str(), val.length());
                        }
                    }
                    RegisterScriptType(typeInfo, metadataKind, effectName);
                }
            }
        } else if (typeInfo->DerivesFrom(itemUseScriptType)) {
            auto metadata = _builder.GetMetadataForType(typeInfo->GetTypeId());
            for (size_t m = 0; m < metadata.size(); m++) {
                auto data = metadata[m];
                if (std::regex_match(data, base_match, metadataMatcher)) {
                    auto mt = base_match[1].str();
                    auto metadataKind = ArbUt::StringView(mt.c_str(), mt.length());
                    auto metadataVariables = base_match[2].str();
                    if (!std::regex_match(metadataVariables, base_match, variableMatcher)) {
                        continue;
                    }
                    ArbUt::StringView effectName;
                    for (size_t variableIndex = 1; variableIndex < base_match.size(); variableIndex += 2) {
                        if (ArbUt::StringView::CalculateHash(base_match[variableIndex].str().c_str()) == "effect"_cnc) {
                            auto val = base_match[variableIndex + 1].str();
                            effectName = ArbUt::StringView(val.c_str(), val.length());
                        }
                    }
                    _itemUseTypes.Insert(effectName, typeInfo);
                }
            }
        }
    }
}
void AngelScriptResolver::RegisterScriptType(asITypeInfo* typeInfo, const ArbUt::StringView& metadataKind,
                                             const ArbUt::StringView& effectName) {
    if (effectName.IsEmpty()) {
        return;
    }
    switch (metadataKind) {
        case "Move"_cnc:
            _typeDatabase[ScriptCategory::Attack].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Pokemon"_cnc:
            _typeDatabase[ScriptCategory::Creature].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Ability"_cnc:
            _typeDatabase[ScriptCategory::Talent].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Status"_cnc:
            _typeDatabase[ScriptCategory::Status].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Battle"_cnc:
            _typeDatabase[ScriptCategory::Battle].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Side"_cnc:
            _typeDatabase[ScriptCategory::Side].Insert(effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
        case "Weather"_cnc:
            _typeDatabase[static_cast<ScriptCategory>(PkmnScriptCategory::Weather)].Insert(
                effectName, new AngelScriptTypeInfo(effectName, typeInfo));
            break;
    }
}
void AngelScriptResolver::CreateScript(const char* name, const char* script) {
    _builder.AddSectionFromMemory(name, script);
}
void AngelScriptResolver::WriteByteCodeToFile(const char* file, bool stripDebugInfo) {
    FILE* wFile = nullptr;
    // Open file in write binary mode.
    wFile = fopen(file, "wb");
    // Ensure we opened it.
    EnsureNotNull(wFile);
    // Save the initial position, as we need to write here later.
    fpos_t startPos;
    fgetpos(wFile, &startPos);
    // The first 8 bytes are the position of the data for which script name is which script decl.
    // We don't know this position yet, so we write an empty value here. We write a size of 8, to accommodate an
    // unsigned long.
    uint64_t byteCodeSizeArr[]{0};
    fwrite(byteCodeSizeArr, sizeof(uint64_t), 1, wFile);
    // We initialize a stream, with no bounds (as this is a read thing)
    auto stream = new FileByteCodeStream(wFile, SIZE_MAX);
    // And save the angelscript byte code.
    _mainModule->SaveByteCode(stream, stripDebugInfo);
    // We grab the current position of the written file. This is the position the types will be written on. So we need
    // to know this.
    uint64_t bytecodeSize = (uint64_t)ftell(wFile);
    stream->WriteTypes(_typeDatabase, _itemUseTypes);

    // Go back to the start of the file
    fsetpos(wFile, &startPos);
    // And write the actual position the types are at.
    byteCodeSizeArr[0] = bytecodeSize;
    fwrite(byteCodeSizeArr, sizeof(uint64_t), 1, wFile);

    // Close the file
    Ensure(fclose(wFile) == 0);
    delete stream;
}
void AngelScriptResolver::LoadByteCodeFromFile(const char* file) {
    FILE* rFile = nullptr;
    // Open the file in read binary mode
    rFile = fopen(file, "rb");
    // Ensure it opened
    EnsureNotNull(rFile);
    // the first 8 bytes are position of the types database, everything before that is the angelscript byte code.
    uint64_t byteCodeSize[]{0};
    fread(byteCodeSize, sizeof(uint64_t), 1, rFile);
    // Initialize a stream, with the earlier found bounds.
    auto stream = new FileByteCodeStream(rFile, byteCodeSize[0] - 1);
    // And load the angelscript byte code.
    int result = _mainModule->LoadByteCode(stream);
    // Ensure we succeeded in this.
    Ensure(result == asSUCCESS);

    // Begin loading the type database
    auto types = stream->ReadTypes();

    InitializeByteCode(types);
    Ensure(fclose(rFile) == 0);
    delete stream;
}
uint8_t* AngelScriptResolver::WriteByteCodeToMemory(size_t& size, bool stripDebugInfo) {
    auto stream = new MemoryByteCodeStream();
    size_t byteCodeSize[]{0};
    stream->Write(byteCodeSize, sizeof(uint64_t));
    auto result = _mainModule->SaveByteCode(stream, stripDebugInfo);
    Ensure(result == asSUCCESS);
    byteCodeSize[0] = (uint64_t)stream->GetWrittenSize();
    stream->WriteTypes(_typeDatabase, _itemUseTypes);
    stream->WriteToPosition(byteCodeSize, sizeof(uint64_t), 0);
    auto arr = stream->GetOut();
    size = stream->GetWrittenSize();
    arr = static_cast<uint8_t*>(realloc(arr, size * sizeof(uint8_t)));
    delete stream;
    return arr;
}
void AngelScriptResolver::LoadByteCodeFromMemory(uint8_t* byte, size_t size) {
    auto stream = new MemoryByteCodeStream(byte, size);
    uint64_t byteCodeSizeArr[]{0};
    stream->Read(byteCodeSizeArr, sizeof(uint64_t));
    stream->SetAngelScriptBound((size_t)byteCodeSizeArr[0]);
    // And load the angelscript byte code.
    int result = _mainModule->LoadByteCode(stream);
    // Ensure we succeeded in this.
    Ensure(result == asSUCCESS);

    // Begin loading the type database
    auto types = stream->ReadTypes();

    InitializeByteCode(types);
    delete stream;
}
void AngelScriptResolver::InitializeByteCode(
    const ArbUt::Dictionary<ScriptCategory, ArbUt::Dictionary<ArbUt::StringView, uint32_t>>& types) {

    auto typeCount = _mainModule->GetObjectTypeCount();
    ArbUt::Dictionary<uint32_t, asITypeInfo*> objectTypes;
    for (asUINT i = 0; i < typeCount; i++) {
        auto t = _mainModule->GetObjectTypeByIndex(i);
        objectTypes.Set(ArbUt::StringView::CalculateHash(t->GetName()), t);
    }
    ArbUt::Dictionary<ScriptCategory, ArbUt::Dictionary<ArbUt::StringView, AngelScriptTypeInfo*>> typeDatabase;
    for (const auto& innerDb : types) {
        if (innerDb.first != (ScriptCategory)-1) {
            ArbUt::Dictionary<ArbUt::StringView, AngelScriptTypeInfo*> newInnerDb;
            for (const auto& val : innerDb.second) {
                auto decl = val.second;
                auto type = objectTypes[decl];
                newInnerDb.Set(val.first, new AngelScriptTypeInfo(val.first, type));
            }
            typeDatabase.Set(innerDb.first, newInnerDb);
        } else {
            for (const auto& val : innerDb.second) {
                auto decl = val.second;
                auto type = objectTypes[decl];
                _itemUseTypes.Set(val.first, type);
            }
        }
    }
    _typeDatabase = typeDatabase;
}