make_package.py

from __future__ import annotations

import argparse
import os
import platform as pf
import shutil
import subprocess
import sys
import time
import zipfile
from pathlib import Path
from loguru import logger

from hoshicore.component.utils import VERSION, SOFTWARE_NAME

join_path = os.path.join

CLI_FILENAME = "launcher"
GUI_FILENAME = "HoshinoWeaver desktop"

logger.remove()
logger.add(sys.stdout, level="INFO")


def _generate_spec(
    work_path: str,
    compile_path: str,
    final_name: str,
    icon_path: str,
    *,
    console_gui: bool = False,
    upx: bool = False,
    macos_bundle: bool = False,
) -> str:
    """Generate a PyInstaller .spec file with MERGE for shared dependencies."""
    gui_script = join_path(work_path, f"{GUI_FILENAME}.py")
    cli_script = join_path(work_path, f"{CLI_FILENAME}.py")

    collect_all_packages = ["pyexiv2"]
    collect_block = "\n".join(
        f"_d, _b, _h = collect_all({pkg!r})\n"
        f"shared_datas += _d; shared_binaries += _b; shared_hiddenimports += _h"
        for pkg in collect_all_packages
    )

    spec = f"""\
# -*- mode: python ; coding: utf-8 -*-
# Auto-generated by make_package.py

from PyInstaller.utils.hooks import collect_all, collect_submodules, copy_metadata

shared_datas = []
shared_binaries = []
shared_hiddenimports = []

{collect_block}

shared_hiddenimports += collect_submodules('scipy._lib.array_api_compat')

# Collect compiled custom-op extension if available
import importlib.util as _imputil
_c_spec = _imputil.find_spec('hoshicore._custom_op._C')
if _c_spec is not None:
    shared_hiddenimports.append('hoshicore._custom_op._C')
    # OpenMP runtime DLL collection:
    #   MSVC  → vcomp140.dll (implicit, collected via ctypes.util)
    #   MinGW → libgomp-1.dll + GCC runtime DLLs copied by build_ops.py into _custom_op/
    try:
        from hoshicore._custom_op import build_info as _build_info
        _bi = _build_info()
        if _bi.get('openmp') and _bi.get('compiler') == 'msvc':
            import ctypes.util as _ctutil
            _vcomp = _ctutil.find_library('vcomp140')
            if _vcomp:
                shared_binaries.append((_vcomp, '.'))
    except Exception:
        pass

    # MinGW builds: build_ops.py copies runtime DLLs next to _C.pyd; collect them all.
    import glob as _glob, os.path as _osp
    _cop_dir = _osp.dirname(_osp.abspath(__import__('hoshicore._custom_op', fromlist=['_custom_op']).__file__))
    for _dll in _glob.glob(_osp.join(_cop_dir, '*.dll')):
        shared_binaries.append((_dll, '.'))

# libturbojpeg: loaded via ctypes at runtime, not traced by PyInstaller
# turbojpeg Python package uses conditional import in image_io, so add explicitly.
try:
    import turbojpeg as _tjpkg
    import ctypes.util as _ctutil_tj, glob as _glob_tj, os.path as _osp_tj
    _tjlib = _ctutil_tj.find_library('turbojpeg')
    if _tjlib:
        shared_binaries.append((_tjlib, '.'))
    else:
        for _f in _glob_tj.glob(_osp_tj.join(_osp_tj.dirname(_tjpkg.__file__), '*turbojpeg*')):
            shared_binaries.append((_f, '.'))
    shared_hiddenimports.append('turbojpeg')
except Exception:
    pass

# pyexiv2 loads exiv2api via sys.path.append + bare import;
# PyInstaller can't trace that, so add the platform extension dir explicitly.
import os as _os, sys as _sys, pyexiv2 as _pyexiv2
_plib = _os.path.join(_os.path.dirname(_pyexiv2.__file__), 'lib')
_pyver = '{{}}.{{}}'.format(_sys.version_info.major, _sys.version_info.minor)
_plat = {{'darwin': 'darwin', 'linux': 'linux', 'win32': 'win'}}[_sys.platform]
#if _plat != 'win':
#    _ext_dir = _os.path.join(_plib, 'py{{}}-{{}}'.format(_pyver, _plat))
#    shared_datas.append((_ext_dir, _os.path.join('pyexiv2', 'lib', 'py{{}}-{{}}'.format(_pyver, _plat))))

# Static resource files (user-modifiable DAG yamls + default settings + LICENSE)
shared_datas.append(({join_path(work_path, 'hoshicore', 'dag')!r}, 'hoshicore/dag'))
shared_datas.append(({join_path(work_path, 'hoshicore', 'default_settings.yaml')!r}, 'hoshicore'))
shared_datas.append(({join_path(work_path, 'LICENSE')!r}, '.'))

_shared_excludes = ['tensorflow', 'keras', 'torch', 'PyQt5', 'PyQt6']

a_gui = Analysis(
    [{gui_script!r}],
    pathex=[{work_path!r}],
    binaries=shared_binaries,
    datas=shared_datas,
    hiddenimports=['PySide6'] + shared_hiddenimports,
    excludes=_shared_excludes,
    noarchive=False,
)

a_cli = Analysis(
    [{cli_script!r}],
    pathex=[{work_path!r}],
    binaries=shared_binaries,
    datas=shared_datas,
    hiddenimports=shared_hiddenimports,
    excludes=['PySide6'] + _shared_excludes,
    noarchive=False,
)

MERGE(
    (a_gui, {GUI_FILENAME!r}, {GUI_FILENAME!r}),
    (a_cli, {CLI_FILENAME!r}, {CLI_FILENAME!r}),
)

pyz_gui = PYZ(a_gui.pure)
exe_gui = EXE(
    pyz_gui,
    a_gui.scripts,
    [],
    exclude_binaries=True,
    name={GUI_FILENAME!r},
    debug=False,
    strip=False,
    upx={upx!r},
    console={console_gui!r},
    icon={icon_path!r},
)

pyz_cli = PYZ(a_cli.pure)
exe_cli = EXE(
    pyz_cli,
    a_cli.scripts,
    [],
    exclude_binaries=True,
    name={CLI_FILENAME!r},
    debug=False,
    strip=False,
    upx={upx!r},
    console=True,
)

coll = COLLECT(
    exe_gui,
    a_gui.binaries,
    a_gui.zipfiles,
    a_gui.datas,
    exe_cli,
    a_cli.binaries,
    a_cli.zipfiles,
    a_cli.datas,
    strip=False,
    upx={upx!r},
    name={final_name!r},
)
"""

    if macos_bundle:
        spec += f"""\

app = BUNDLE(
    coll,
    name={GUI_FILENAME + '.app'!r},
    icon={icon_path!r},
    bundle_identifier='com.hoshinoweaver.desktop',
)
"""

    return spec


def file_to_zip(path_original: str, z: zipfile.ZipFile) -> None:
    f_list = list(Path(path_original).glob("**/*"))
    for f in f_list:
        z.write(f, str(f)[len(path_original):])


platform_mapping = {
    "win32": "win",
    "cygwin": "win",
    "darwin": "macos",
    "linux2": "linux",
    "linux": "linux",
}


def _ensure_custom_ops(work_path: str, *, skip_build: bool, allow_numpy_only: bool) -> bool:
    """Verify _C is importable; build if needed. Return True if available."""
    result = subprocess.run(
        [sys.executable, "-c", "import hoshicore._custom_op._C"],
        capture_output=True,
    )
    if result.returncode == 0:
        logger.info("Custom ops (_C) verified.")
        return True
    if skip_build:
        if allow_numpy_only:
            logger.warning("Custom ops (_C) not available. Package will use numpy fallback.")
            return False
        logger.error(
            "Custom ops (_C) not available. "
            "Build first with `python csrc/build_ops.py`, or pass --allow-numpy-only."
        )
        sys.exit(1)
    logger.info("Custom ops (_C) not found, attempting build...")
    ret = subprocess.run(
        [sys.executable, join_path(work_path, "csrc", "build_ops.py")],
        cwd=work_path,
    )
    if ret.returncode != 0:
        if allow_numpy_only:
            logger.warning("Custom-op build failed. Continuing with numpy fallback.")
            return False
        logger.error("Custom-op build failed. Fix build errors or pass --allow-numpy-only.")
        sys.exit(1)
    logger.info("Custom ops built successfully.")
    return True


def main():
    argparser = argparse.ArgumentParser(
        description="Package HoshiNoWeaver using PyInstaller.")
    argparser.add_argument(
        "--apply-upx",
        action="store_true",
        help="Apply UPX to squeeze the size of the executable.")
    argparser.add_argument(
        "--apply-zip",
        action="store_true",
        help="Generate .zip file after packaging.")
    argparser.add_argument(
        "--debug-gui",
        action="store_true",
        help="Generate GUI version with console window.")
    argparser.add_argument(
        "--no-build",
        action="store_true",
        help="Skip automatic C++ custom-op compilation (assume already built or absent).")
    argparser.add_argument(
        "--allow-numpy-only",
        action="store_true",
        help="Allow packaging without compiled _C (numpy fallback). Default: error if _C unavailable.")
    args = argparser.parse_args()

    platform = platform_mapping[sys.platform]
    exec_suffix = ""
    if platform == "win":
        exec_suffix = ".exe"
    if platform == "macos":
        mac_main_ver = int(pf.mac_ver()[0].split(".")[0])
        if mac_main_ver >= 13:
            platform += "13+"

    work_path = os.path.dirname(os.path.abspath(__file__))
    compile_path = join_path(work_path, "dist")
    build_path = join_path(work_path, "build")

    t0 = time.time()

    FINAL_DIR_NAME = (
        f"{GUI_FILENAME}_{platform}_{VERSION}"
        f"{'-debug' if args.debug_gui else ''}")

    icon_path = join_path(work_path, "imgs", "HNW.jpg")
    macos_bundle = platform.startswith("macos") and not args.debug_gui

    # ── Ensure custom ops ──
    _ensure_custom_ops(
        work_path,
        skip_build=args.no_build,
        allow_numpy_only=args.allow_numpy_only,
    )

    # ── Generate .spec ──
    spec_content = _generate_spec(
        work_path,
        compile_path,
        final_name=FINAL_DIR_NAME,
        icon_path=icon_path,
        console_gui=args.debug_gui,
        upx=args.apply_upx,
        macos_bundle=macos_bundle,
    )

    spec_path = join_path(build_path, "HoshiNoWeaver.spec")
    os.makedirs(build_path, exist_ok=True)
    with open(spec_path, "w", encoding="utf-8") as f:
        f.write(spec_content)
    logger.info(f"Generated spec file: {spec_path}")

    # ── Run PyInstaller ──
    cmd = [
        sys.executable, "-m", "PyInstaller",
        "--clean",
        "--distpath", compile_path,
        "--workpath", build_path,
        spec_path,
    ]
    logger.info(f"Running: {' '.join(cmd)}")
    ret_code = subprocess.run(cmd).returncode
    if ret_code != 0:
        logger.error("PyInstaller build failed.")
        sys.exit(-1)
    logger.info("Build completed successfully.")

    # ── ZIP ──
    if args.apply_zip:
        zip_fname = join_path(compile_path, f"{SOFTWARE_NAME}_{platform}_{VERSION}.zip")
        logger.info(f"Zipping to {zip_fname} ...")
        with zipfile.ZipFile(zip_fname, mode="w",
                             compression=zipfile.ZIP_DEFLATED) as zf:
            file_to_zip(join_path(compile_path, FINAL_DIR_NAME), zf)
        logger.info("Zip done.")

    logger.info(
        f"Package script finished. Total time cost {(time.time()-t0):.2f}s.")


if __name__ == "__main__":
    main()