#include "payload/payload.h"
#include "common/timers/Timer.h"
#include "core/serialize/serialize.h"
#include "payload/binary/binary.h"
#include "payload/text/text.h"

using namespace std;

namespace sorbet::payload {

void createInitialGlobalState(core::GlobalState &gs, const realmain::options::Options &options,
                              const unique_ptr<const OwnedKeyValueStore> &kvstore) {
    if (options.cacheSensitiveOptions.noStdlib) {
        gs.initEmpty();
        return;
    }

    if (PAYLOAD_EMPTY) {
        Timer timeit(gs.tracer(), "read_global_state.source");
        sorbet::rbi::populateRBIsInto(gs);
    } else {
        Timer timeit(gs.tracer(), "read_global_state.binary");
        core::serialize::Serializer::loadGlobalState(gs, PAYLOAD_SYMBOL_TABLE, PAYLOAD_NAME_TABLE, PAYLOAD_FILE_TABLE);
    }
    ENFORCE(gs.utf8NamesUsed() < core::GlobalState::PAYLOAD_MAX_UTF8_NAME_COUNT,
            "Payload defined `{}` UTF8 names, which is greater than the expected maximum of `{}`. Consider updating "
            "`PAYLOAD_MAX_UTF8_NAME_COUNT` in `GlobalState`.",
            gs.utf8NamesUsed(), core::GlobalState::PAYLOAD_MAX_UTF8_NAME_COUNT);
    ENFORCE(
        gs.constantNamesUsed() < core::GlobalState::PAYLOAD_MAX_CONSTANT_NAME_COUNT,
        "Payload defined `{}` Constant names, which is greater than the expected maximum of `{}`. Consider updating "
        "`PAYLOAD_MAX_CONSTANT_NAME_COUNT` in `GlobalState`.",
        gs.constantNamesUsed(), core::GlobalState::PAYLOAD_MAX_CONSTANT_NAME_COUNT);
    ENFORCE(gs.uniqueNamesUsed() < core::GlobalState::PAYLOAD_MAX_UNIQUE_NAME_COUNT,
            "Payload defined `{}` Unique names, which is greater than the expected maximum of `{}`. Consider updating "
            "`PAYLOAD_MAX_UNIQUE_NAME_COUNT` in `GlobalState`.",
            gs.uniqueNamesUsed(), core::GlobalState::PAYLOAD_MAX_UNIQUE_NAME_COUNT);
    ENFORCE(gs.fieldsUsed() < core::GlobalState::PAYLOAD_MAX_FIELD_COUNT,
            "Payload defined `{}` fields, which is greater than the expected maximum of `{}`. Consider updating "
            "`PAYLOAD_MAX_FIELD_COUNT` in `GlobalState`.",
            gs.fieldsUsed(), core::GlobalState::PAYLOAD_MAX_FIELD_COUNT);
    ENFORCE(gs.methodsUsed() < core::GlobalState::PAYLOAD_MAX_METHOD_COUNT,
            "Payload defined `{}` methods, which is greater than the expected maximum of `{}`. Consider updating "
            "`PAYLOAD_MAX_METHOD_COUNT` in `GlobalState`.",
            gs.methodsUsed(), core::GlobalState::PAYLOAD_MAX_METHOD_COUNT);
    ENFORCE(gs.classAndModulesUsed() < core::GlobalState::PAYLOAD_MAX_CLASS_AND_MODULE_COUNT,
            "Payload defined `{}` classes and modules, which is greater than the expected maximum of `{}`. Consider "
            "updating `PAYLOAD_MAX_CLASS_AND_MODULE_COUNT` in `GlobalState`.",
            gs.classAndModulesUsed(), core::GlobalState::PAYLOAD_MAX_CLASS_AND_MODULE_COUNT);
    ENFORCE(gs.typeMembersUsed() < core::GlobalState::PAYLOAD_MAX_TYPE_MEMBER_COUNT,
            "Payload defined `{}` type members, which is greater than the expected maximum of `{}`. Consider updating "
            "`PAYLOAD_MAX_TYPE_MEMBER_COUNT` in `GlobalState`.",
            gs.typeMembersUsed(), core::GlobalState::PAYLOAD_MAX_TYPE_MEMBER_COUNT);
    ENFORCE(
        gs.typeParametersUsed() < core::GlobalState::PAYLOAD_MAX_TYPE_ARGUMENT_COUNT,
        "Payload defined `{}` type arguments, which is greater than the expected maximum of `{}`. Consider updating "
        "`PAYLOAD_MAX_TYPE_ARGUMENT_COUNT` in `GlobalState`.",
        gs.typeParametersUsed(), core::GlobalState::PAYLOAD_MAX_TYPE_ARGUMENT_COUNT);

    // We can use the kvstore to read in the cached name table. We read this in after the payload has been initialized,
    // as the cached name table will extend the payload's existing table when the sorbet versions match.

    // We've just read the payload from the cache, so it would make sense to call gs.markNameTableAsCached() here. That
    // way, we wouldn't to store the payload to the cache again in the name table cache for the first stratum. However,
    // when we start readingthe cache entries for non payload names, we'll resize the names hash, which makes the
    // entries added from the payload invalid. To get around this, we don't mark the name table as cached, and then just
    // store the names from the payload again.

    if (kvstore) {
        auto maybeUUIDBytes = kvstore->read(core::serialize::Serializer::NAME_TABLE_UUID_KEY);
        auto maybeDiffCountAndHashSizeBytes =
            kvstore->read(core::serialize::Serializer::NAME_TABLE_DIFF_COUNT_AND_HASH_SIZE_KEY);
        if (maybeUUIDBytes.data != nullptr) {
            Timer timeit(gs.tracer(), "read_name_table.kvstore");
            if (maybeDiffCountAndHashSizeBytes.data != nullptr) {
                auto diffCount = core::serialize::Serializer::loadNameTableDiffCountAndResizeNamesHash(
                    gs, maybeDiffCountAndHashSizeBytes.data);
                ENFORCE(diffCount > 0);

                // Load UUID
                ENFORCE(gs.kvstoreUuid == 0, "The name table may only be loaded into a fresh GlobalState");
                gs.kvstoreUuid = core::serialize::Serializer::loadGlobalStateUUID(gs, maybeUUIDBytes.data);

                // TODO(neil): In the future, we could compact diffs to avoid reading many small entries.
                for (uint32_t i = 0; i < diffCount; i++) {
                    auto data = kvstore->read(core::serialize::Serializer::nameTableDiffKey(i));
                    ENFORCE(data.data != nullptr);
                    core::serialize::Serializer::loadAndAppendNameTableDiff(gs, data.data);
                }
                gs.setNameTableDiffCount(diffCount);
                gs.markNameTableAsCached();
                // Check that well-known names are consistent with what the payload would expect.
                if constexpr (debug_mode) {
                    // The name of the Sorbet class is well-known, and the name is generated by
                    // core/tools/generate_names.cc. If this doesn't match up, something has gone seriously wrong
                    // with name table serialization.
                    ENFORCE(core::Symbols::Sorbet().data(gs)->name == core::Names::Constants::Sorbet());
                    ENFORCE(core::Names::Constants::Sorbet().kind() == core::NameKind::CONSTANT);
                    auto name = core::Symbols::Sorbet().data(gs)->name.dataCnst(gs)->original;
                    ENFORCE(name.kind() == core::NameKind::UTF8);
                    ENFORCE(name.dataUtf8(gs)->utf8 == "Sorbet");

                    // Array#bsearch is defined in the payload rbis with no overloads. If that changes, pick a new
                    // method that's defined in an rbi, but on a well-known class.
                    ENFORCE(core::Symbols::T_Array().data(gs)->name.kind() == core::NameKind::CONSTANT);
                    auto bsearch = gs.lookupMethodSymbol(core::Symbols::Array(), gs.lookupNameUTF8("bsearch"));
                    ENFORCE(bsearch.exists());
                    name = bsearch.data(gs)->name;
                    ENFORCE(name.kind() == core::NameKind::UTF8);
                    ENFORCE(name.dataUtf8(gs)->utf8 == "bsearch");
                }
            }
            if constexpr (debug_mode) {
                for (unsigned int i = 1; i < gs.filesUsed(); i++) {
                    core::FileRef fref(i);

                    // We should only see payload files at this point.
                    ENFORCE(fref.dataAllowingUnsafe(gs).sourceType != core::File::Type::Normal);
                }
            }
        }
    }
}

} // namespace sorbet::payload
