# Copyright 2024 The JAX Authors.
#
# 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
#
#     https://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.

"""Setup script for JAX ROCm plugin package."""

import importlib
import os
import sys
from setuptools import setup
from setuptools.dist import Distribution

# The PEP 517 backend does not put the source root on sys.path.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rocm_version import DEFAULT_ROCM_PATH, detect_rocm_version

__version__ = None
rocm_version = 0  # placeholder
project_name = f"jax-rocm{rocm_version}-plugin"
package_name = f"jax_rocm{rocm_version}_plugin"

# Hermetic wheel actions expose WHEEL_VERSION_SUFFIX, but not ROCM_PATH.
rocm_path = os.getenv("ROCM_PATH", DEFAULT_ROCM_PATH)
rocm_tag = os.getenv("ROCM_VERSION_EXTRA")
rocm_detected_version = detect_rocm_version(rocm_path, rocm_tag)

def load_version_module(pkg_path):
  spec = importlib.util.spec_from_file_location(
    'version', os.path.join(pkg_path, 'version.py'))
  module = importlib.util.module_from_spec(spec)
  spec.loader.exec_module(module)
  return module

_version_module = load_version_module(package_name)
__version__ = _version_module._get_version_for_build()
if rocm_tag:
  __version__ = __version__ + "+rocm" + rocm_tag
_cmdclass = _version_module._get_cmdclass(package_name)

class BinaryDistribution(Distribution):
  """This class makes 'bdist_wheel' include an ABI tag on the wheel."""

  def has_ext_modules(self):
    return True

_description = f"JAX Plugin for AMD GPUs (ROCm:{rocm_detected_version})"

setup(
    name=project_name,
    version=__version__,
    cmdclass=_cmdclass,
    description=_description,
    long_description=_description,
    long_description_content_type="text/plain",
    author="ROCm JAX Devs",
    author_email="dl.dl-JAX@amd.com",
    packages=[package_name],
    python_requires=">=3.12",
    install_requires=[f"jax-rocm{rocm_version}-pjrt=={__version__}"],
    url="https://github.com/jax-ml/jax",
    license="Apache-2.0",
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Programming Language :: Python :: 3.12",
        "Programming Language :: Python :: 3.13",
        "Programming Language :: Python :: 3.14",
    ],
    package_data={
        package_name: [
            "*",
        ],
    },
    zip_safe=False,
    distclass=BinaryDistribution,
    extras_require={
        "with_rocm": [
            "amd_rocm_hip_runtime_devel_instinct",
            "amd_rocm_hip_runtime_instinct",
            "amd_hipblas_instinct",
            "amd_hipsparse_instinct",
            "amd_hipsolver_instinct",
            "amd_miopen_hip_instinct",
            "amd_rocm_llvm_instinct",
            "amd_rocm_language_runtime_instinct",
            "amd_rccl_instinct",
            "amd_hipfft_instinct",
            "amd_rocm_device_libs_instinct",
            "amd_hipsparselt_instinct",
        ],
    },
)