# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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
#
# http://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.
import os
import platform
from pathlib import Path
from typing import List

from setuptools import find_packages, setup
from setuptools.dist import Distribution


def parse_requirements(filename: os.PathLike):
    with open(filename) as f:
        requirements = f.read().splitlines()

        def extract_url(line):
            return next(filter(lambda x: x[0] != '-', line.split()))

        extra_URLs = []
        deps = []
        for line in requirements:
            if line.startswith("#") or line.startswith("-r") or line.startswith(
                    "-c"):
                continue

            # handle -i and --extra-index-url options
            if "-i " in line or "--extra-index-url" in line:
                extra_URLs.append(extract_url(line))
            # handle URLs such as git+https://github.com/flashinfer-ai/flashinfer.git@e3853dd#egg=flashinfer-python
            elif line.startswith("git+https"):
                idx = line.find("egg=")
                dep = line[idx + 4:]
                deps.append(dep)
            else:
                deps.append(line)
    return deps, extra_URLs


def sanity_check():
    tensorrt_llm_path = Path(__file__).resolve().parent / "tensorrt_llm"
    if not ((tensorrt_llm_path / "bindings").exists() or
            (tensorrt_llm_path / "bindings.pyi").exists()):
        raise ImportError(
            'The `bindings` module does not exist. Please check the package integrity. '
            'If you are attempting to use the pip development mode (editable installation), '
            'please execute `scripts/build_wheel.py` first, and then run `pip install -e .`.'
        )


def get_version():
    version_file = Path(
        __file__).resolve().parent / "tensorrt_llm" / "version.py"
    version = None
    with open(version_file) as f:
        for line in f:
            if not line.startswith("__version__"):
                continue
            version = line.split('"')[1]

    if version is None:
        raise RuntimeError(f"Could not set version from {version_file}")

    return version


class BinaryDistribution(Distribution):
    """Distribution which always forces a binary package with platform name"""

    def has_ext_modules(self):
        return True


on_windows = platform.system() == "Windows"
required_deps, extra_URLs = parse_requirements(
    Path("requirements-windows.txt" if on_windows else "requirements.txt"))
devel_deps, _ = parse_requirements(
    Path("requirements-dev-windows.txt"
         if on_windows else "requirements-dev.txt"))
constraints_file = Path("constraints.txt")
if constraints_file.exists():
    constraints, _ = parse_requirements(constraints_file)
    required_deps.extend(constraints)

if on_windows:
    package_data = [
        'libs/th_common.dll', 'libs/tensorrt_llm.dll',
        'libs/nvinfer_plugin_tensorrt_llm.dll', 'bindings.*.pyd', "include/**/*"
    ]
else:
    package_data = [
        'bin/executorWorker', 'libs/libtensorrt_llm.so', 'libs/libth_common.so',
        'libs/libnvinfer_plugin_tensorrt_llm.so',
        'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so',
        'libs/libtensorrt_llm_nixl_wrapper.so',
        'libs/libdecoder_attention_1.so', 'libs/nvshmem/License.txt',
        'libs/nvshmem/nvshmem_bootstrap_uid.so.3',
        'libs/nvshmem/nvshmem_transport_ibgda.so.103', 'bindings.*.so',
        'deep_ep/LICENSE', 'deep_ep_cpp_tllm.*.so', "include/**/*"
    ]

package_data += [
    'bindings/*.pyi',
    'tools/plugin_gen/templates/*',
    'bench/build/benchmark_config.yml',
    'evaluate/lm_eval_tasks/**/*',
    "_torch/auto_deploy/config/*.yaml",
]


def download_precompiled(workspace: str) -> str:
    import glob
    import subprocess

    from setuptools.errors import SetupError

    cmd = [
        "pip", "download", f"tensorrt_llm={get_version()}",
        f"--dest={workspace}", "--no-deps",
        "--extra-index-url=https://pypi.nvidia.com"
    ]
    try:
        subprocess.check_call(cmd)
        wheel_path = glob.glob(f"{workspace}/tensorrt_llm-*.whl")[0]
    except Exception as e:
        raise SetupError(
            "Failed to download the automatically resolved wheel, please try specifying TRTLLM_USE_PRECOMPILED with a link or local path to a valid wheel."
        ) from e
    else:
        return wheel_path


def extract_from_precompiled(precompiled_location: str, package_data: List[str],
                             workspace: str) -> None:
    """Extract package data (binaries and other materials) from a precompiled wheel to the working directory.
    This allows skipping the compilation, and repackaging the binaries and Python files in the working directory to a new wheel.
    """
    import fnmatch
    import tarfile
    import zipfile
    from urllib.request import urlretrieve

    from setuptools.errors import SetupError

    if os.path.isfile(precompiled_location):
        precompiled_path = precompiled_location
        print(f"Using local precompiled file: {precompiled_path}.")
    else:
        precompiled_filename = precompiled_location.split("/")[-1]
        precompiled_path = os.path.join(workspace, precompiled_filename)
        print(
            f"Downloading precompiled file from {precompiled_location} to {precompiled_path}."
        )
        try:
            urlretrieve(precompiled_location, filename=precompiled_path)
        except Exception as e:
            raise SetupError(
                f"Failed to get precompiled file from {precompiled_location}."
            ) from e

    if precompiled_path.endswith("tar.gz"):
        with tarfile.open(precompiled_path, "r:gz") as tar:
            for member in tar.getmembers():
                if fnmatch.fnmatchcase(member.name,
                                       "TensorRT-LLM/tensorrt_llm-*.whl"):
                    break
            else:
                raise SetupError(
                    f"Failed to get wheel file from {precompiled_path}.") from e

            wheel_path = os.path.join(workspace, member.name)
            tar.extract(member, path=workspace, filter=tarfile.data_filter)
    else:
        wheel_path = precompiled_path

    with zipfile.ZipFile(wheel_path) as wheel:
        for file in wheel.filelist:
            if file.filename.endswith((".py", ".yaml")):
                continue
            for filename_pattern in package_data:
                if fnmatch.fnmatchcase(file.filename,
                                       f"tensorrt_llm/{filename_pattern}"):
                    break
            else:
                continue
            print(
                f"Extracting and including {file.filename} from precompiled wheel."
            )
            wheel.extract(file)


use_precompiled: bool = os.getenv("TRTLLM_USE_PRECOMPILED") == "1"
precompiled_location: str = os.getenv("TRTLLM_PRECOMPILED_LOCATION")

if precompiled_location:
    use_precompiled = True

if use_precompiled:
    from tempfile import TemporaryDirectory
    with TemporaryDirectory() as tempdir:
        if not precompiled_location:
            precompiled_location = download_precompiled(tempdir)
        extract_from_precompiled(precompiled_location, package_data, tempdir)

sanity_check()

# https://setuptools.pypa.io/en/latest/references/keywords.html
setup(
    name='tensorrt_llm',
    version=get_version(),
    description='TensorRT-LLM: A TensorRT Toolbox for Large Language Models',
    long_description=
    'TensorRT-LLM: A TensorRT Toolbox for Large Language Models',
    author="NVIDIA Corporation",
    url="https://github.com/NVIDIA/TensorRT-LLM",
    download_url="https://github.com/NVIDIA/TensorRT-LLM/tags",
    packages=find_packages(),
    # TODO Add windows support for python bindings.
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Developers",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.12",
    ],
    distclass=BinaryDistribution,
    license="Apache License 2.0",
    keywords="nvidia tensorrt deeplearning inference",
    package_data={
        'tensorrt_llm': package_data,
    },
    entry_points={
        'console_scripts': [
            'trtllm-build=tensorrt_llm.commands.build:main',
            'trtllm-prune=tensorrt_llm.commands.prune:main',
            'trtllm-refit=tensorrt_llm.commands.refit:main',
            'trtllm-bench=tensorrt_llm.commands.bench:main',
            'trtllm-serve=tensorrt_llm.commands.serve:main',
            'trtllm-eval=tensorrt_llm.commands.eval:main'
        ],
    },
    scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'],
    extras_require={
        "devel": devel_deps,
    },
    zip_safe=True,
    install_requires=required_deps,
    dependency_links=
    extra_URLs,  # Warning: Dependency links support has been dropped by pip 19.0
    python_requires=">=3.7, <4")