import importlib.util
import logging
import os
import re
import subprocess
import sys
from pathlib import Path
from shutil import which
from typing import Dict, List

from packaging.version import Version, parse
from setuptools import Extension, find_packages, setup
from setuptools.command.build_ext import build_ext


def load_module_from_path(module_name, path):
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


ROOT_DIR = os.path.dirname(__file__)
logger = logging.getLogger(__name__)

# cannot import envs directly because it depends on vllm,
#  which is not installed yet
envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'vllm', 'envs.py'))

VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE

if not sys.platform.startswith("linux"):
    logger.warning(
        "vLLM only supports Linux platform (including WSL). "
        "Building on %s, "
        "so vLLM may not be able to run correctly", sys.platform)
    VLLM_TARGET_DEVICE = "empty"


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


def find_version(filepath: str, version_name: str = '__version__') -> str:
    """Extract version information from the given filepath.

    Adapted from https://github.com/ray-project/ray/blob/0b190ee1160eeca9796bc091e07eaebf4c85b511/python/setup.py
    """
    with open(filepath) as fp:
        version_match = re.search(r"^{} = ['\"]([^'\"]*)['\"]".format(version_name),
                                  fp.read(), re.M)
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


def get_commit_id() -> str:
    """
    get the current commit of vllm mlu
    """
    git_short_hash = subprocess.run(
        ['git', 'rev-parse', '--short', 'HEAD'],
        stdout=subprocess.PIPE,
        text=True
    ).stdout.strip()
    return git_short_hash


def get_vllm_version() -> str:
    """
    get vllm version
    """
    vllm_version_path = get_path("vllm", "version.py")
    version = (
        find_version(vllm_version_path, '__version__')
        + "+mlu" + find_version(vllm_version_path, '__vllm_mlu_version__')
        + ".pt" + find_version(vllm_version_path, '__torch_version__')
    )
    if get_commit_id():
        version += "." + get_commit_id()
    return version


def read_readme() -> str:
    """Read the README file if present."""
    p = get_path("README.md")
    if os.path.isfile(p):
        with open(get_path("README.md"), encoding="utf-8") as f:
            return f.read()
    else:
        return ""


def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
            requirements = f.read().strip().split("\n")
        resolved_requirements = []
        for line in requirements:
            if line.startswith("-r "):
                resolved_requirements += _read_requirements(line.split()[1])
            elif line.startswith("--"):
                continue
            else:
                resolved_requirements.append(line)
        return resolved_requirements

    requirements = _read_requirements("requirements-mlu.txt")
    return requirements

package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json", "version_config"]
}
if envs.VLLM_USE_PRECOMPILED:
    ext_modules = []
    package_data["vllm"].append("*.so")


setup(
    name="vllm",
    version=get_vllm_version(),
    author="Cambricon vLLM Team",
    license="Apache 2.0",
    description=("A high-throughput and memory-efficient inference and "
                 "serving engine for LLMs on MLU backendon"),
    long_description=read_readme(),
    long_description_content_type="text/markdown",
    url="",
    project_urls={
        "Homepage": "https://github.com/vllm-project/vllm",
        "Documentation": "https://vllm.readthedocs.io/en/latest/",
    },
    classifiers=[
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: Python :: 3.12",
        "License :: OSI Approved :: Apache Software License",
        "Intended Audience :: Developers",
        "Intended Audience :: Information Technology",
        "Intended Audience :: Science/Research",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Scientific/Engineering :: Information Analysis",
    ],
    packages=find_packages(exclude=("benchmarks", "docs", "examples", "tests*")),
    python_requires=">=3.8",
    install_requires=get_requirements(),
    extras_require={
        "tensorizer": ["tensorizer>=2.9.0"],
        "audio": ["librosa", "soundfile"],  # Required for audio processing
        "video": ["decord"]  # Required for video processing
    },
    package_data=package_data,
    entry_points={
        "console_scripts": [
            "vllm=vllm.scripts:main",
        ],
    },
)