#     Copyright 2026, Kay Hayen, mailto:kay.hayen@gmail.com find license text at end of file


"""Standard plug-in to make multiprocessing and joblib work well.

On Windows, the multiprocessing modules forks new processes which then have
to start from scratch. This won't work if there is no "sys.executable" to
point to a "Python.exe" and won't use compiled code by default.

The issue applies to accelerated and standalone mode alike.

spell-checker: ignore joblib,anyio
"""

from nuitka.importing.FakeModules import makeFakeModuleDescription
from nuitka.ModuleRegistry import (
    getModuleInclusionInfoByName,
    getRootTopModule,
)
from nuitka.options.Options import isStandaloneMode, shallMakeModule
from nuitka.plugins.PluginBase import NuitkaPluginBase
from nuitka.PythonVersions import python_version
from nuitka.tree.SourceHandling import readSourceCodeFromFilename
from nuitka.utils.ModuleNames import ModuleName


class NuitkaPluginMultiprocessingWorkarounds(NuitkaPluginBase):
    """This is to make multiprocessing work with Nuitka and use compiled code.

    When running in accelerated mode, it's not good to fork a new Python
    instance to run other code, as that won't be accelerated. And when
    run in standalone mode, there may not even be a Python, but it's the
    same principle.

    So by default, this module is on and works around the behavior of the
    "multiprocessing.forking/multiprocessing.spawn/multiprocessing.manager"
    expectations.
    """

    plugin_name = "multiprocessing"
    plugin_desc = "Required by 'multiprocessing' package."
    plugin_category = "package-support"

    def __init__(self):
        self.created_parents_main = False

    @classmethod
    def isRelevant(cls):
        return not shallMakeModule()

    @staticmethod
    def isAlwaysEnabled():
        return True

    @staticmethod
    def createPreModuleLoadCode(module):
        full_name = module.getFullName()

        if full_name == "__parents_main__":
            yield (
                """\
import sys
# If the multiprocessing fork flag is present, scrub the arguments so that the
# user code (which might use argparse) doesn't see them. Save them in builtins
# so we can restore them later when calling multiprocessing.spawn.spawn_main.
if "--multiprocessing-fork" in sys.argv:
    # Use builtins to store the args so they are available globally even in functions
    # without needing to pass them around or use global keyword everywhere.
    if str is bytes:
        import __builtin__ as builtins
    else:
        import builtins

    builtins.__nuitka_original_args = sys.argv[:]
    sys.argv = [sys.argv[0]]
""",
                "Scrubbing multiprocessing fork arguments from sys.argv.",
            )

        if full_name in ("multiprocessing", "anyio"):
            yield (
                """\
import sys, os
argv0 = sys.argv[0]
if sys.platform == "win32" and not os.path.exists(argv0) and not argv0.endswith(".exe"):
    argv0 += ".exe"

sys.executable = %s
sys._base_executable = sys.executable
""" % ("__nuitka_binary_exe" if isStandaloneMode() else "argv0"),
                """\
Monkey patching "%s" load environment.""" % full_name,
            )

    @staticmethod
    def createPostModuleLoadCode(module):
        full_name = module.getFullName()

        if full_name == "multiprocessing":
            yield (
                """\
try:
    from multiprocessing.forking import ForkingPickler
except ImportError:
    from multiprocessing.reduction import ForkingPickler

import os
import sys

class C:
   def f():
       pass

def _reduce_compiled_method(m):
    if m.im_self is None:
        return getattr, (m.im_class, m.im_func.__name__)
    else:
        return getattr, (m.im_self, m.im_func.__name__)

ForkingPickler.register(type(C().f), _reduce_compiled_method)
if str is bytes:
    ForkingPickler.register(type(C.f), _reduce_compiled_method)

if str is not bytes:
    import multiprocessing.spawn

    if not hasattr(multiprocessing.spawn, "__nuitka_original_fixup_main_from_path"):
        multiprocessing.spawn.__nuitka_original_fixup_main_from_path = multiprocessing.spawn._fixup_main_from_path

        def _fixup_main_from_path_for_nuitka(main_path):
            assert main_path is not None

            import __parents_main__ as parents_main

            sys.modules["__main__"] = parents_main
            sys.modules["__mp_main__"] = parents_main

        multiprocessing.spawn._fixup_main_from_path = _fixup_main_from_path_for_nuitka
""",
                """\
Monkey patching "multiprocessing" for compiled methods.""",
            )

        if full_name == "anyio.to_process":
            yield (
                """\
import anyio.to_process as _anyio_to_process
import os
import pickle
import sys
from importlib.util import module_from_spec, spec_from_file_location

def _process_worker_for_nuitka():
    stdin = sys.stdin
    stdout = sys.stdout
    sys.stdin = open(os.devnull)
    sys.stdout = open(os.devnull, "w")

    stdout.buffer.write(b"READY\\n")
    stdout.buffer.flush()

    while True:
        retval = exception = None

        try:
            payload = pickle.load(stdin.buffer)
            command = payload[0]
            args = payload[1:]
        except EOFError:
            return
        except BaseException as exc:
            exception = exc
        else:
            if command == "run":
                func = args[0]
                func_args = args[1]

                try:
                    retval = func(*func_args)
                except BaseException as exc:
                    exception = exc
            elif command == "init":
                sys.path, main_module_path = args

                del sys.modules["__main__"]

                if main_module_path and os.path.isfile(main_module_path):
                    try:
                        spec = spec_from_file_location("__mp_main__", main_module_path)

                        if spec and spec.loader:
                            main = module_from_spec(spec)
                            spec.loader.exec_module(main)
                            sys.modules["__main__"] = main
                            sys.modules["__mp_main__"] = main
                    except BaseException as exc:
                        exception = exc
                else:
                    try:
                        import __parents_main__ as parents_main
                    except ImportError:
                        pass
                    else:
                        sys.modules["__main__"] = parents_main
                        sys.modules["__mp_main__"] = parents_main

        try:
            if exception is not None:
                status = b"EXCEPTION"
                pickled = pickle.dumps(exception, pickle.HIGHEST_PROTOCOL)
            else:
                status = b"RETURN"
                pickled = pickle.dumps(retval, pickle.HIGHEST_PROTOCOL)
        except BaseException as exc:
            exception = exc
            status = b"EXCEPTION"
            pickled = pickle.dumps(exc, pickle.HIGHEST_PROTOCOL)

        stdout.buffer.write(b"%s %d\\n" % (status, len(pickled)))
        stdout.buffer.write(pickled)
        stdout.buffer.flush()

        if isinstance(exception, SystemExit):
            raise exception

_anyio_to_process.process_worker = _process_worker_for_nuitka
""",
                """\
Monkey patching "anyio.to_process" for worker start.""",
            )

    def createFakeModuleDependency(self, module):
        full_name = module.getFullName()

        if full_name not in ("multiprocessing", "anyio"):
            return

        if (
            self.created_parents_main
            or getModuleInclusionInfoByName("__parents_main__") is not None
        ):
            return

        # First, build the module node and then read again from the
        # source code.
        root_module = getRootTopModule()

        # TODO: Might have to also make "__mp_main__"
        module_name = ModuleName("__parents_main__")

        source_code = readSourceCodeFromFilename(module_name, root_module.getFilename())

        # For the call stack, this may look bad or different to what CPython
        # does, but such is life. The forking module was split up into multiple
        # modules in Python3 so this is doubled. TODO: Move the version check
        # into the code and let Nuitka compile time optimize it, which will also
        # make sure joblib support is there for Python2 as well.
        if python_version >= 0x300:
            source_code += """
def __nuitka_freeze_support():
    import sys
    import builtins
    import multiprocessing.spawn

    # This is a variant of freeze_support that will work for multiprocessing and
    # joblib equally well.
    kwds = {}

    if not hasattr(builtins, "__nuitka_original_args"):
        return

    sys.argv = builtins.__nuitka_original_args
    del builtins.__nuitka_original_args

    for arg in sys.argv[2:]:
        try:
            name, value = arg.split('=')
        except ValueError:
            name = "pipe_handle"
            value = arg

        if value == 'None':
            kwds[name] = None
        else:
            kwds[name] = int(value)

    # Otherwise main module names will not work.
    sys.modules["__main__"] = sys.modules["__parents_main__"]

    multiprocessing.spawn.spawn_main(**kwds)
__nuitka_freeze_support()
"""
        else:
            source_code += """
def __nuitka_freeze_support():
    import __builtin__, sys
    if not hasattr(__builtin__, "__nuitka_original_args"):
        return

    sys.argv = __builtin__.__nuitka_original_args
    del __builtin__.__nuitka_original_args

    sys.modules["__main__"] = sys.modules[__name__]
    __import__("multiprocessing.forking").forking.freeze_support()
__nuitka_freeze_support()
"""

        yield makeFakeModuleDescription(
            module_name=module_name,
            source_code=source_code,
            source_filename=root_module.getCompileTimeFilename(),
            reason="Auto enable multiprocessing/anyio freeze support",
        )

        self.created_parents_main = True

    def onModuleEncounter(
        self, using_module_name, module_name, module_filename, module_kind
    ):
        # Enforce recursion in to multiprocessing for accelerated mode, which
        # would normally avoid this.
        if module_name.hasNamespace("multiprocessing"):
            return True, "Multiprocessing plugin needs this to monkey patch it."

    def decideCompilation(self, module_name):
        if module_name.hasNamespace("multiprocessing"):
            return "bytecode"

        # TODO: Make this demoted too.
        # or module_name in( "multiprocessing-preLoad", "multiprocessing-postLoad"):

    @staticmethod
    def getPreprocessorSymbols():
        if getModuleInclusionInfoByName("__parents_main__"):
            return {"_NUITKA_PLUGIN_MULTIPROCESSING_ENABLED": "1"}


#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the GNU Affero General Public License, Version 3 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        https://www.gnu.org/licenses/agpl-3.0.txt
#
#     See also: "Nuitka Runtime Library Exception, Version 1.0" in file
#     "LICENSE-RUNTIME.txt" for additional permissions granted under Section 7.
#
#     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.
