# SPDX-License-Identifier: Apache-2.0
# Derived from vLLM; modified by DeepSeek.

import os

import setuptools
import torch
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension

from deep_cluster.hardware import accelerator_kind

IS_NPU = accelerator_kind() == "ascend"

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

# Extension configuration shared by the CUDA and NPU backends.
COMMON_COMPILE_FLAGS = ["-O2", "-std=c++20", "-Wno-deprecated-declarations"]
ABI_FLAG = "-D_GLIBCXX_USE_CXX11_ABI=1"
assert torch._C._GLIBCXX_USE_CXX11_ABI, "must use CXX11 ABI"

CXX_FLAGS = ["-g1", *COMMON_COMPILE_FLAGS, ABI_FLAG]
COMMON_SOURCES = [
    "csrc/pybind.cpp",
    "csrc/engine/pybind.cpp",
    "csrc/engine/pybind_sequence.cpp",
    "csrc/engine/pybind_scheduler.cpp",
    "csrc/engine/pybind_output.cpp",
    "csrc/engine/zos.cpp",
    "csrc/engine/input_metadata/block_load_plan.cpp",
    "csrc/engine/input_metadata/prepare_tensor_input_base.cpp",
    "csrc/engine/input_metadata/prepare_tensor_input_v4.cpp",
    "csrc/engine/prepare_sample.cpp",
    "csrc/graph_timer/pybind.cpp",
]
COMMON_LIBRARIES = ["pthread"]
INCLUDE_DIRS = [os.path.abspath("csrc")]

if IS_NPU:
    from torch_npu.utils.cpp_extension import NpuExtension  # pyright: ignore[reportMissingImports]

    deepinfer_extension = NpuExtension(
        name="deepinfer._C",
        sources=[
            *COMMON_SOURCES,
            "csrc/asc_utils.cpp",
            "csrc/graph_timer/graph_timer_asc.cpp",
        ],
        extra_compile_args={
            "cxx": [*CXX_FLAGS, "-fopenmp", "-fvisibility=hidden", "-DWITH_NPU"],
        },
        libraries=COMMON_LIBRARIES,
        include_dirs=INCLUDE_DIRS,
    )
else:
    conda_prefix = os.environ.get("CONDA_PREFIX") or os.environ.get("VIRTUAL_ENV")
    if not conda_prefix:
        raise RuntimeError("Neither CONDA_PREFIX nor VIRTUAL_ENV is set; a conda env or virtualenv is required to build.")

    deepinfer_extension = CUDAExtension(
        name="deepinfer._C",
        sources=[
            *COMMON_SOURCES,
            "csrc/cuda_utils.cpp",
            "csrc/graph_timer/graph_timer_cuda.cpp",
        ],
        extra_compile_args={
            "cxx": [*CXX_FLAGS, "-fvisibility=hidden"],
        },
        libraries=[*COMMON_LIBRARIES, "cuda"],
        library_dirs=[os.path.join(CUDA_HOME, "lib64", "stubs")],  # for libcuda.so
        include_dirs=INCLUDE_DIRS,
    )


def get_path(*filepath: str) -> str:
    return os.path.join(ROOT_DIR, *filepath)


def get_requirements() -> list[str]:
    """Get Python package dependencies from requirements.txt (comments and blank lines stripped)."""
    with open(get_path("requirements.txt")) as f:
        lines = (line.split("#", 1)[0].strip() for line in f)
        return [line for line in lines if line]


setuptools.setup(
    packages=setuptools.find_packages(include=("deepinfer", "deepinfer.*", "deep_cluster", "deep_cluster.*")),
    install_requires=get_requirements(),
    ext_modules=[deepinfer_extension],
    cmdclass={"build_ext": BuildExtension},
)