//     Copyright 2019, Kay Hayen, mailto:kay.hayen@gmail.com
//
//     Part of "Nuitka", an optimizing Python compiler that is compatible and
//     integrates with CPython, but also works on its own.
//
//     Licensed under the Apache License, Version 2.0 (the "License");
//     you may not use this file except in compliance with the License.
//     You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//     Unless required by applicable law or agreed to in writing, software
//     distributed under the License is distributed on an "AS IS" BASIS,
//     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//     See the License for the specific language governing permissions and
//     limitations under the License.
//
/* The main program for a compiled program.
 *
 * It needs to prepare the interpreter and then loads and executes
 * the "__main__" module.
 * *
 */

#include "nuitka/prelude.h"

#include "build_definitions.h"

#include "osdefs.h"
#include "structseq.h"

#if defined(_WIN32)
#include <windows.h>
#endif

extern PyCodeObject *codeobj_main;
extern PyObject *const_str_plain___main__;

/* For later use in "Py_GetArgcArgv" */
static char **orig_argv;
static int orig_argc;
#if PYTHON_VERSION >= 300
static wchar_t **argv_unicode;
#endif

#ifdef _NUITKA_STANDALONE

#if _NUITKA_FROZEN > 0
extern void copyFrozenModulesTo(struct _frozen *destination);
#endif

static void prepareStandaloneEnvironment() {
    // Tell the CPython library to use our pre-compiled modules as frozen
    // modules. This for those modules/packages like "encoding" that will be
    // loaded during "Py_Initialize" already, for the others they may be
    // compiled.

#if _NUITKA_FROZEN > 0
    // The CPython library has some pre-existing frozen modules, we only append
    // to that.
    struct _frozen const *search = PyImport_FrozenModules;
    while (search->name) {
        search++;
    }
    int pre_existing_count = (int)(search - PyImport_FrozenModules);

    /* Allocate new memory and merge the tables. Keeping the old ones has
     * the advantage that e.g. "import this" is going to work well.
     */
    struct _frozen *merged =
        (struct _frozen *)malloc(sizeof(struct _frozen) * (_NUITKA_FROZEN + pre_existing_count + 1));

    memcpy(merged, PyImport_FrozenModules, pre_existing_count * sizeof(struct _frozen));
    copyFrozenModulesTo(merged + pre_existing_count);
    PyImport_FrozenModules = merged;
#endif

    /* Setup environment variables to tell CPython that we would like it to use
     * the provided binary directory as the place to look for DLLs and for
     * extension modules.
     */
#if defined(_WIN32)
    SetDllDirectoryW(getBinaryDirectoryWideChars());
#endif

#if PYTHON_VERSION < 300
    char *binary_directory = getBinaryDirectoryHostEncoded();
    NUITKA_PRINTF_TRACE("Binary dir is %s\n", binary_directory);

    Py_SetPythonHome(binary_directory);
#else
    wchar_t *binary_directory = getBinaryDirectoryWideChars();
    NUITKA_PRINTF_TRACE("Binary dir is %S\n", binary_directory);

    Py_SetPythonHome(binary_directory);

#endif
}

#if PYTHON_VERSION < 300
#define PY_FORMAT_GETPATH_RESULT "%s"
#else
#define PY_FORMAT_GETPATH_RESULT "%ls"
#endif

static void restoreStandaloneEnvironment() {
    /* Make sure to use the optimal value for standalone mode only. */
#if PYTHON_VERSION < 300
    PySys_SetPath(getBinaryDirectoryHostEncoded());
    NUITKA_PRINTF_TRACE("Final PySys_GetPath is 's'.\n", PySys_GetPath());
#else
    PySys_SetPath(getBinaryDirectoryWideChars());
    Py_SetPath(getBinaryDirectoryWideChars());
    NUITKA_PRINTF_TRACE("Final Py_GetPath is '%ls'.\n", Py_GetPath());
#endif
}

#endif

extern void _initCompiledCellType();
extern void _initCompiledGeneratorType();
extern void _initCompiledFunctionType();
extern void _initCompiledMethodType();
extern void _initCompiledFrameType();
#if PYTHON_VERSION >= 350
extern void _initCompiledCoroutineTypes();
#endif
#if PYTHON_VERSION >= 360
extern void _initCompiledAsyncgenTypes();
#endif

#include <locale.h>

// Types of command line arguments are different between Python2/3.
#if PYTHON_VERSION >= 300
typedef wchar_t **argv_type_t;
static argv_type_t convertCommandLineParameters(int argc, char **argv) {
#if _WIN32
    int new_argc;

    argv_type_t result = CommandLineToArgvW(GetCommandLineW(), &new_argc);
    assert(new_argc == argc);
    return result;
#else
    // Originally taken from CPython3: There seems to be no sane way to use
    static wchar_t **argv_copy;
    argv_copy = (wchar_t **)PyMem_Malloc(sizeof(wchar_t *) * argc);

    // Temporarily disable locale for conversions to not use it.
    char *oldloc = strdup(setlocale(LC_ALL, NULL));
    setlocale(LC_ALL, "");

    for (int i = 0; i < argc; i++) {
#if PYTHON_VERSION >= 350
        argv_copy[i] = Py_DecodeLocale(argv[i], NULL);
#elif defined(__APPLE__) && PYTHON_VERSION >= 320
        argv_copy[i] = _Py_DecodeUTF8_surrogateescape(argv[i], strlen(argv[i]));
#else
        argv_copy[i] = _Py_char2wchar(argv[i], NULL);
#endif

        assert(argv_copy[i]);
    }

    setlocale(LC_ALL, oldloc);
    free(oldloc);

    return argv_copy;
#endif
}
#else
typedef char **argv_type_t;
#endif

#ifdef _NUITKA_PLUGIN_WINDOWS_SERVICE_ENABLED
extern void SvcInstall();
extern void SvcLaunchService();

// Callback from Windows Service logic.
DWORD WINAPI SvcStartPython(LPVOID lpParam) {
    IMPORT_EMBEDDED_MODULE(const_str_plain___main__, "__main__");

    return 0;
}
#endif

// Parse the command line parameters and provide it to "sys" built-in module,
// as well as decide if it's a multiprocessing usage.

static bool setCommandLineParameters(int argc, argv_type_t argv, bool initial) {
    bool is_multiprocessing_fork = false;

    if (initial) {
        /* We might need to handle special parameters from plugins that are
           very deeply woven into command line handling. These are right now
           multiprocessing, which indicates that it's forking via extra
           command line argument. And Windows Service indicates need to
           install and exit here too.
         */
        for (int i = 1; i < argc; i++) {
#if PYTHON_VERSION < 300
            if ((strcmp(argv[i], "--multiprocessing-fork")) == 0 && (i + 1 < argc))
#else
            // TODO: Should simply use wide char literal
            wchar_t constant_buffer[100];
            mbstowcs(constant_buffer, "--multiprocessing-fork", 100);
            if ((wcscmp(argv[i], constant_buffer)) == 0 && (i + 1 < argc))
#endif
            {
                is_multiprocessing_fork = true;
                break;
            }

#ifdef _NUITKA_PLUGIN_WINDOWS_SERVICE_ENABLED
#if PYTHON_VERSION < 300
            if ((strcmp(argv[i], "install")) == 0 && (i + 1 < argc))
#else
            // TODO: Should simply use wide char literal
            wchar_t constant_buffer[100];
            mbstowcs(constant_buffer, "install", 100);
            if ((wcscmp(argv[i], constant_buffer)) == 0 && (i + 1 < argc))
#endif
            {
                SvcInstall();
                NUITKA_CANNOT_GET_HERE("SvcInstall must not return");
            }
#endif
        }
    }

    if (initial) {
        // Py_SetProgramName(argv[0]);
    } else {
        PySys_SetArgv(argc, argv);
    }

    return is_multiprocessing_fork;
}

#ifdef _NUITKA_WINMAIN_ENTRY_POINT
int __stdcall WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, char *lpCmdLine, int nCmdShow) {
#if defined(__MINGW32__) && !defined(_W64)
    /* MINGW32 */
    int argc = _argc;
    char **argv = _argv;
#else
    /* MSVC, MINGW64 */
    int argc = __argc;
    char **argv = __argv;
#endif
#else
int main(int argc, char **argv) {
#endif
    NUITKA_PRINT_TRACE("main(): Entered.");

    orig_argv = argv;
    orig_argc = argc;

#ifdef __FreeBSD__
    /* 754 requires that FP exceptions run in "no stop" mode by default, and
     *
     *    until C vendors implement C99's ways to control FP exceptions, Python
     *    requires non-stop mode.  Alas, some platforms enable FP exceptions by
     *    default.  Here we disable them.
     */

    fp_except_t m;

    m = fpgetmask();
    fpsetmask(m & ~FP_X_OFL);
#endif

    /* On Windows we support loading the constants blob from an embedded
     * resource. On Linux, where possible this is done automatically by
     * the linker already.
     */
#if defined(_NUITKA_CONSTANTS_FROM_RESOURCE)
    NUITKA_PRINT_TRACE("main(): Loading constants blob from Windows resource.");

    loadConstantsResource();
#endif

#ifdef _NUITKA_STANDALONE
    NUITKA_PRINT_TRACE("main(): Prepare standalone environment.");
    prepareStandaloneEnvironment();
#else

#if PYTHON_VERSION >= 350 && defined(DLL_EXTRA_PATH)
    NUITKA_PRINT_TRACE("main(): Prepare DLL extra path.");
    SetDllDirectory(DLL_EXTRA_PATH);
#endif

#endif

    /* Initialize CPython library environment. */
    Py_DebugFlag = 0;
#if PYTHON_VERSION < 300
    Py_Py3kWarningFlag = _NUITKA_SYSFLAG_PY3K_WARNING;
    Py_DivisionWarningFlag = _NUITKA_SYSFLAG_DIVISION_WARNING;
    Py_UnicodeFlag = _NUITKA_SYSFLAG_UNICODE;
    Py_TabcheckFlag = 0;
#endif
    Py_InspectFlag = 0;
    Py_InteractiveFlag = 0;
    Py_OptimizeFlag = _NUITKA_SYSFLAG_OPTIMIZE;
    Py_DontWriteBytecodeFlag = 0;
    Py_NoUserSiteDirectory = _NUITKA_SYSFLAG_NO_SITE;
    Py_IgnoreEnvironmentFlag = 0;
    Py_VerboseFlag = _NUITKA_SYSFLAG_VERBOSE;
    Py_BytesWarningFlag = _NUITKA_SYSFLAG_BYTES_WARNING;
#if _NUITKA_SYSFLAG_BYTES_WARNING
    Py_HashRandomizationFlag = 1;
#endif
#if PYTHON_VERSION >= 370
    Py_UTF8Mode = _NUITKA_SYSFLAG_UTF8;
#endif

    /* This suppresses warnings from getpath.c */
    Py_FrozenFlag = 1;

    /* We want to import the site module, but only after we finished our own
     * setup. The site module import will be the first thing, the main module
     * does.
     */
    Py_NoSiteFlag = 1;

    /* Initial command line handling only. */

#if PYTHON_VERSION >= 300
    NUITKA_PRINT_TRACE("main(): Calling convertCommandLineParameters.");
    argv_unicode = convertCommandLineParameters(argc, argv);
#endif

    NUITKA_PRINT_TRACE("main(): Calling setCommandLineParameters.");

#if PYTHON_VERSION < 300
    bool is_multiprocess_forking = setCommandLineParameters(argc, argv, true);
#else
    bool is_multiprocess_forking = setCommandLineParameters(argc, argv_unicode, true);
#endif

    /* For Python installations that need the home set, we inject it back here. */
#if defined(PYTHON_HOME_PATH)
#if PYTHON_VERSION < 300
    NUITKA_PRINT_TRACE("main(): Prepare run environment '" PYTHON_HOME_PATH "'.");
    Py_SetPythonHome(PYTHON_HOME_PATH);
#else
    NUITKA_PRINTF_TRACE("main(): Prepare run environment '%S'.\n", L"" PYTHON_HOME_PATH);
    Py_SetPythonHome(L"" PYTHON_HOME_PATH);
    // Make sure the above Py_SetPythonHome call has effect already.
    Py_GetPath();
#endif
#endif

    /* Initialize the embedded CPython interpreter. */
    NUITKA_PRINT_TRACE("main(): Calling Py_Initialize to initialize interpreter.");
    Py_Initialize();

#ifdef _NUITKA_STANDALONE
    NUITKA_PRINT_TRACE("main(): Restore standalone environment.");
    restoreStandaloneEnvironment();
#endif

    /* Lie about it, believe it or not, there are "site" files, that check
     * against later imports, see below.
     */
    Py_NoSiteFlag = _NUITKA_SYSFLAG_NO_SITE;

    /* Set the command line parameters for run time usage. */
    NUITKA_PRINT_TRACE("main(): Calling setCommandLineParameters.");

#if PYTHON_VERSION < 300
    setCommandLineParameters(argc, argv, false);
#else
    setCommandLineParameters(argc, argv_unicode, false);
#endif

    /* Initialize the built-in module tricks used. */
    NUITKA_PRINT_TRACE("main(): Calling _initBuiltinModule().");
    _initBuiltinModule();

    /* Initialize the Python constant values used. This also sets
     * "sys.executable" while at it.
     */
    NUITKA_PRINT_TRACE("main(): Calling createGlobalConstants().");
    createGlobalConstants();

    NUITKA_PRINT_TRACE("main(): Calling _initBuiltinOriginalValues().");
    _initBuiltinOriginalValues();

    /* Revert the wrong "sys.flags" value, it's used by "site" on at least
     * Debian for Python 3.3, more uses may exist.
     */
#if _NUITKA_SYSFLAG_NO_SITE == 0
#if PYTHON_VERSION < 300
    PyStructSequence_SET_ITEM(PySys_GetObject((char *)"flags"), 9, const_int_0);
#else
    PyStructSequence_SetItem(PySys_GetObject("flags"), 6, const_int_0);
#endif
#endif

    /* Initialize the compiled types of Nuitka. */
    _initCompiledCellType();
    _initCompiledGeneratorType();
    _initCompiledFunctionType();
    _initCompiledMethodType();
    _initCompiledFrameType();
#if PYTHON_VERSION >= 350
    _initCompiledCoroutineTypes();
#endif
#if PYTHON_VERSION >= 360
    _initCompiledAsyncgenTypes();
#endif

#if PYTHON_VERSION < 300
    _initSlotCompare();
#endif
#if PYTHON_VERSION >= 270
    _initSlotIternext();
#endif

    NUITKA_PRINT_TRACE("main(): Calling enhancePythonTypes().");
    enhancePythonTypes();

    NUITKA_PRINT_TRACE("main(): Calling patchBuiltinModule().");
    patchBuiltinModule();

    NUITKA_PRINT_TRACE("main(): Calling patchTypeComparison().");
    patchTypeComparison();

    NUITKA_PRINT_TRACE("main(): Calling patchTracebackDealloc().");
    patchTracebackDealloc();

    /* Allow to override the ticker value, to remove checks for threads in
     * CPython core from impact on benchmarks. */
    char const *ticker_value = getenv("NUITKA_TICKER");
    if (ticker_value != NULL) {
        _Py_Ticker = atoi(ticker_value);
        assert(_Py_Ticker >= 20);
    }

#ifdef _NUITKA_STANDALONE
    NUITKA_PRINT_TRACE("main(): Calling setEarlyFrozenModulesFileAttribute().");

#if PYTHON_VERSION >= 300
    PyObject *os_module = PyImport_ImportModule("os");
    CHECK_OBJECT(os_module);
#endif
    setEarlyFrozenModulesFileAttribute();
#endif

    NUITKA_PRINT_TRACE("main(): Calling setupMetaPathBasedLoader().");
    /* Enable meta path based loader. */
    setupMetaPathBasedLoader();

    _PyWarnings_Init();

    /* Disable CPython warnings if requested to. */
#if _NUITKA_NO_PYTHON_WARNINGS
    {
#if PYTHON_VERSION >= 300
        wchar_t ignore[] = L"ignore";
#else
        char ignore[] = "ignore";
#endif

        PySys_AddWarnOption(ignore);

#if PYTHON_VERSION >= 342 && defined(_NUITKA_FULL_COMPAT)
        // For full compatibility bump the warnings registry version,
        // otherwise modules "__warningsregistry__" will mismatch.
        PyObject *warnings_module = PyImport_ImportModule("warnings");
        PyObject *meth = PyObject_GetAttrString(warnings_module, "_filters_mutated");

        CALL_FUNCTION_NO_ARGS(meth);
#endif
    }
#endif

#if PYTHON_VERSION >= 300
    NUITKA_PRINT_TRACE("main(): Calling patchInspectModule().");
    patchInspectModule();
#endif

#if _NUITKA_PROFILE
    startProfiling();
#endif

    /* Execute the main module. In case of multiprocessing making a fork on
     * Windows, we should execute something else instead. */
    if (unlikely(is_multiprocess_forking)) {
        NUITKA_PRINT_TRACE("main(): Calling __parents_main__.");
        IMPORT_EMBEDDED_MODULE(PyUnicode_FromString("__parents_main__"), "__parents_main__");
    } else {
        PyDict_DelItem(PyImport_GetModuleDict(), const_str_plain___main__);

#ifndef _NUITKA_PLUGIN_WINDOWS_SERVICE_ENABLED
        /* Execute the "__main__" module. */
        NUITKA_PRINT_TRACE("main(): Calling __main__.");

        IMPORT_EMBEDDED_MODULE(const_str_plain___main__, "__main__");
#else
        SvcLaunchService();
#endif
    }

#if _NUITKA_PROFILE
    stopProfiling();
#endif

#ifndef __NUITKA_NO_ASSERT__
    checkGlobalConstants();

    /* TODO: Walk over all loaded compiled modules, and make this kind of checks. */
#if 0
    checkModuleConstants___main__();
#endif

#endif

    if (ERROR_OCCURRED()) {
#if PYTHON_VERSION >= 300
        /* Remove the frozen importlib traceback part, which would not be compatible. */
        PyThreadState *thread_state = PyThreadState_GET();

        while (thread_state->curexc_traceback) {
            PyTracebackObject *tb = (PyTracebackObject *)thread_state->curexc_traceback;
            PyFrameObject *frame = tb->tb_frame;

            if (0 == strcmp(PyUnicode_AsUTF8(frame->f_code->co_filename), "<frozen importlib._bootstrap>")) {
                thread_state->curexc_traceback = (PyObject *)tb->tb_next;
                Py_INCREF(tb->tb_next);

                continue;
            }

            break;
        }
#endif

        PyErr_PrintEx(0);
        Py_Exit(1);
    } else {
        Py_Exit(0);
    }

    /* The above branches both do "Py_Exit()" calls which are not supposed to
     * return.
     */
    NUITKA_CANNOT_GET_HERE(main);
}

/* This is an unofficial API, not available on Windows, but on Linux and others
 * it is exported, and has been used by some code.
 */
#ifndef _WIN32
#ifdef __cplusplus
extern "C" {
#endif

#if PYTHON_VERSION >= 300
#if defined(__GNUC__)
__attribute__((visibility("default")))
#endif
void Py_GetArgcArgv(int *argc, wchar_t ***argv) {
    *argc = orig_argc;
    *argv = argv_unicode;
}

#else
#if defined(__GNUC__)
__attribute__((visibility("default")))
#endif
void Py_GetArgcArgv( int *argc, char ***argv ) {
    *argc = orig_argc;
    *argv = orig_argv;
}
#endif

#ifdef __cplusplus
}
#endif

#endif
