# ------------------------------------------------------------------
# Copyright (c) 2023 PyInstaller Development Team.
#
# This file is distributed under the terms of the GNU General Public
# License (version 2.0 or later).
#
# The full license is available in LICENSE, distributed with
# this software.
#
# SPDX-License-Identifier: GPL-2.0-or-later
# ---------------------------------------------------

from PyInstaller.utils.hooks import collect_data_files, collect_dynamic_libs, collect_submodules, is_module_satisfies

hiddenimports = []
datas = []

# Ensure that triton/_C/libtriton.so is collected
binaries = collect_dynamic_libs('triton')

# triton has a JIT module that requires its source .py files. For some god-forsaken reason, this JIT module
# (`triton.runtime.jit` attempts to directly read the contents of file pointed to by its `__file__` attribute (assuming
# it is a source file). Therefore, `triton.runtime.jit` must not be collected into PYZ. Same goes for `compiler` and
# `language` sub-packages.
module_collection_mode = {
    'triton': 'pyz+py',
    'triton.runtime.jit': 'py',
    'triton.compiler': 'py',
    'triton.language': 'py',
}

# triton 3.0.0 introduced `triton.backends` sub-package with backend-specific files.
if is_module_satisfies('triton >= 3.0.0'):
    # Collect backend sub-modules/packages.
    hiddenimports += collect_submodules('triton.backends')

    # At the time of writing (triton v3.1.0), `triton.backends.amd` is a namespace package, and is not captured by the
    # above `collect_submodules` call.
    hiddenimports += collect_submodules('triton.backends.amd')

    # Collect ptxas compiler files from `triton/backends/nvidia`, and the HIP/ROCm files from `triton/backends/amd`.
    datas += collect_data_files('triton.backends')
else:
    # Collect ptxas compiler files from triton/third_party/cuda directory. Strictly speaking, the ptxas executable from
    # bin directory should be collected as a binary, but in this case, it makes no difference (plus, PyInstaller >= 6.0
    # has automatic binary-vs-data reclassification).
    datas += collect_data_files('triton.third_party.cuda')