#!/usr/bin/env python3

"""
Create bundle with preloaded modules.

Usage:

    py-preloaded-bundle-fork-server.py <module> <module> ... -o <output-runtime-file>

    <output-runtime-file> <python-script>

Example:

    # original command:
    python my_script.py --foo bar --baz 123

    # create bundle:
    py-preloaded-bundle-fork-server.py pytorch -o python-pytorch.bin

    # run with bundle, faster startup:
    ./python-pytorch.bin my_script.py --foo bar --baz 123
"""

from __future__ import annotations
from typing import BinaryIO
import argparse
import importlib
import sys
import os
import textwrap
from preloaded import criu, startup, _io


def main():
    """main entry"""
    arg_parser = argparse.ArgumentParser(
        epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    arg_parser.add_argument("modules", nargs="+", help="modules to preload")
    arg_parser.add_argument("-o", "--output", required=True, help="output runtime file")
    args = arg_parser.parse_args()

    with open(args.output, "w") as f:
        f.write(textwrap.dedent(f"""\
            #!/usr/bin/env python3
            {'"' * 3}
            This file is auto-generated by py-preloaded-bundle-fork-server.py
            Preloaded modules: {args.modules}
            {'"' * 3}

            import os
            from preloaded.startup import startup_via_fork_server

            if __name__ == "__main__":
                startup_via_fork_server(modules={args.modules!r})
            """))
    os.chmod(args.output, 0o755)

    print("Dumped executable state, created startup helper script:", args.output)


if __name__ == '__main__':
    main()
