import io
import os
import re
import shutil
import subprocess
import sys
import argparse
from typing import Any, Optional

from distutils.core import Command
from setuptools import Extension, find_namespace_packages, setup
from setuptools.command.build_ext import build_ext

try:
    from setuptools.command.bdist_wheel import bdist_wheel
except ModuleNotFoundError:
    from wheel.bdist_wheel import bdist_wheel

from scripts.build_support.env import (
    get_cxx_abi,
    set_cuda_envs,
    set_ilu_envs,
    set_mlu_envs,
    set_musa_envs,
    set_npu_envs,
)
from scripts.build_support.utils import (
    check_and_install_pre_commit,
    get_base_dir,
    get_cmake_dir,
    get_cpu_arch,
    get_device_type,
    get_python_version,
    get_torch_version,
    get_version,
    pre_build,
    read_readme,
)

BUILD_TEST_FILE: bool = True
BUILD_EXPORT: bool = True


def _maybe_compile_tilelang_kernels(device: str) -> None:
    if device != "npu":
        return

    output_root = os.path.join(get_cmake_dir(), "xllm", "compiler", "tilelang")
    os.makedirs(output_root, exist_ok=True)

    env = os.environ.copy()
    base_dir = get_base_dir()

    cmd = [
        sys.executable,
        os.path.join(base_dir, "xllm", "compiler", "tilelang_launcher.py"),
        "compile-kernels",
        "--target",
        "ascend",
        "--output-root",
        output_root,
    ]
    print("[INFO] compiling TileLang kernels via source-tree launcher")
    subprocess.check_call(cmd, cwd=base_dir, env=env)
        
class CMakeExtension(Extension):
    def __init__(self, name: str, path: str, sourcedir: str = "") -> None:
        super().__init__(name, sources=[])
        self.sourcedir = os.path.realpath(os.path.abspath(sourcedir))
        self.path = path

class ExtBuild(build_ext):
    user_options = build_ext.user_options + [
        ("base-dir=", None, "base directory of xLLM project"),
        ("device=", None, "target device type (npu or mlu or cuda or ilu or musa)"),
        ("arch=", None, "target arch type (x86 or arm)"),
        ("generate-so=", None, "generate so or binary"),
    ]

    def initialize_options(self) -> None:
        build_ext.initialize_options(self)
        self.base_dir = get_base_dir()
        self.device: Optional[str] = None
        self.arch: Optional[str] = None
        self.generate_so: bool = False

    def finalize_options(self) -> None:
        build_ext.finalize_options(self)

    def run(self) -> None:
        # check if cmake is installed
        try:
            out: bytes = subprocess.check_output(["cmake", "--version"])
        except OSError:
            raise RuntimeError(
                "CMake must be installed to build the following extensions: "
                + ", ".join(e.name for e in self.extensions)
            )
            exit(1)

        match = re.search(
            r"version\s*(?P<major>\d+)\.(?P<minor>\d+)([\d.]+)?", out.decode()
        )
        if match is None:
            raise RuntimeError(f"Failed to parse CMake version from: {out!r}")
        cmake_major, cmake_minor = int(match.group("major")), int(match.group("minor"))
        if (cmake_major, cmake_minor) < (3, 18):
            raise RuntimeError("CMake >= 3.18.0 is required")

        try:
            # build extensions
            for ext in self.extensions:
                self.build_extension(ext)
        except Exception as e:
            print("ERROR: Build failed.")
            print(f"Details: {e}")
            exit(1)

    def build_extension(self, ext: CMakeExtension) -> None:
        ninja_dir = shutil.which("ninja")
        # the output dir for the extension
        extdir: str = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))

        # create build directory
        os.makedirs(self.build_temp, exist_ok=True)

        # Using this requires trailing slash for auto-detection & inclusion of
        # auxiliary "native" libs

        debug: int = int(os.environ.get("DEBUG", 0)) if self.debug is None else int(self.debug)
        build_type: str = "Debug" if debug else "Release"

        default_jobs = os.cpu_count() or 1
        max_jobs: str = os.getenv("MAX_JOBS", str(default_jobs))
        max_jobs_int: int = int(max_jobs)
        
        # Limit archive (ar/ranlib) concurrency to avoid file locking conflicts.
        # The ar tool requires exclusive access to archive files (.a files) when
        # creating or updating static libraries. When multiple ar processes attempt
        # to modify the same archive file simultaneously, they compete for file locks,
        # which can cause deadlocks and hang the build process.
        archive_jobs: int = min(8, max(1, max_jobs_int // 4))
        cmake_args: list[str] = [
            "-G",
            "Ninja",
            f"-DCMAKE_MAKE_PROGRAM={ninja_dir}",
            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
            f"-DCMAKE_RUNTIME_OUTPUT_DIRECTORY={extdir}",
            "-DUSE_CCACHE=ON",
            f"-DPython_EXECUTABLE:FILEPATH={sys.executable}",
            f"-DCMAKE_BUILD_TYPE={build_type}",
            f"-DBUILD_SHARED_LIBS=OFF",
            f"-DDEVICE_TYPE=USE_{self.device.upper()}",
            f"-DDEVICE_ARCH={self.arch.upper()}",
            f"-DXLLM_ATB_LAYERS_SOURCE_DIR={os.path.join(self.base_dir, 'third_party', 'xllm_atb_layers')}",
            f"-DCMAKE_JOB_POOLS=archive={archive_jobs}",
        ]

        if self.device is None:
            raise ValueError("Please set --device to npu or mlu or cuda or ilu or musa.")
        if self.arch is None:
            raise ValueError("Please set --arch to x86 or arm.")

        if self.device == "npu":
            cmake_args += ["-DUSE_NPU=ON"]
            set_npu_envs()
            _maybe_compile_tilelang_kernels(self.device)
        elif self.device == "mlu":
            cmake_args += ["-DUSE_MLU=ON"]
            set_mlu_envs()
        elif self.device == "cuda":
            torch_cuda_architectures = os.getenv("TORCH_CUDA_ARCH_LIST")
            if not torch_cuda_architectures:
                raise ValueError("Please set TORCH_CUDA_ARCH_LIST environment variable, e.g. export TORCH_CUDA_ARCH_LIST=\"8.0 8.9 9.0 10.0 12.0\"")
            cmake_args += ["-DUSE_CUDA=ON", 
                           f"-DTORCH_CUDA_ARCH_LIST={torch_cuda_architectures}"]
            set_cuda_envs()
        elif self.device == "ilu":
            cmake_args += ["-DUSE_ILU=ON"]
            set_ilu_envs()
        elif self.device == "musa":
            cmake_args += ["-DUSE_MUSA=ON"]
            set_musa_envs()
            global BUILD_TEST_FILE
            BUILD_TEST_FILE = False
        else:
            raise ValueError("Please set --device to npu or mlu or cuda or ilu or musa.")

        product: str = "xllm"
        if self.generate_so:
            product = "libxllm.so"
            cmake_args += ["-DGENERATE_SO=ON"]
        else:
            cmake_args += ["-DGENERATE_SO=OFF"]

        # Adding CMake arguments set as environment variable
        # (needed e.g. to build for ARM OSx on conda-forge)
        if "CMAKE_ARGS" in os.environ:
            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]

        # check if torch binary is built with cxx11 abi
        if get_cxx_abi():
            cmake_args += ["-DUSE_CXX11_ABI=ON", "-D_GLIBCXX_USE_CXX11_ABI=1"]
        else:
            cmake_args += ["-DUSE_CXX11_ABI=OFF", "-D_GLIBCXX_USE_CXX11_ABI=0"]
        
        build_args = ["--config", build_type]
        build_args += ["-j" + max_jobs]

        env: dict[str, str] = os.environ.copy()
        env["VCPKG_MAX_CONCURRENCY"] = str(max_jobs)
        print("CMake Args: ", cmake_args)
        print("Env: ", env)

        self.build_cmake_targets(ext, cmake_args, build_args, env, extdir, product)

    def build_cmake_targets(
        self,
        ext: CMakeExtension,
        cmake_args: list[str],
        build_args: list[str],
        env: dict[str, str],
        extdir: str,
        product: str,
    ) -> None:
        """Build CMake targets"""
        cmake_dir = get_cmake_dir()
        subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)

        base_build_args = build_args
        # add build target to speed up the build process
        build_args += ["--target", ext.name, "xllm"]
        subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)

        os.makedirs(os.path.join(os.path.dirname(cmake_dir), "xllm/core/server/"), exist_ok=True)
        shutil.copy(
            os.path.join(extdir, product),
            os.path.join(os.path.dirname(cmake_dir), "xllm/core/server/"),
        )

        if BUILD_EXPORT:
            # build export module
            build_args = base_build_args + ["--target export_module"]
            subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)

        if BUILD_TEST_FILE:
            # build tests target
            build_args = base_build_args + ["--target all_tests"]
            subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)

class ExtBuildSingleTest(ExtBuild):
    """Inherit ExtBuild, used to build and run a single test"""
    user_options = ExtBuild.user_options + [
        ("test-name=", None, "name of the test target to build and run"),
    ]

    def initialize_options(self) -> None:
        ExtBuild.initialize_options(self)
        self.test_name: Optional[str] = None

    def finalize_options(self) -> None:
        ExtBuild.finalize_options(self)
        if not self.test_name:
            raise ValueError("--test-name is required for ExtBuildSingleTest")

    def build_cmake_targets(
        self,
        ext: CMakeExtension,
        cmake_args: list[str],
        build_args: list[str],
        env: dict[str, str],
        extdir: str,
        product: str,
    ) -> None:
        """Override method: only build the specified test target and run"""
        cmake_dir = get_cmake_dir()
        subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)

        base_build_args = build_args
        # Only build the specified test target
        build_args += ["--target", self.test_name]
        subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)

        # Find test executable
        # CMake usually places executables in CMAKE_RUNTIME_OUTPUT_DIRECTORY or build directory
        test_executable: Optional[str] = None
        possible_paths: list[str] = [
            os.path.join(cmake_dir, self.test_name),
            os.path.join(extdir, self.test_name),
            os.path.join(cmake_dir, "xllm", "core", self.test_name),
        ]
        
        # Check possible paths first
        for path in possible_paths:
            if os.path.exists(path) and os.access(path, os.X_OK):
                test_executable = path
                break
        
        # If not found, try recursive search in build directory
        if not test_executable:
            for root, dirs, files in os.walk(cmake_dir):
                if self.test_name in files:
                    candidate = os.path.join(root, self.test_name)
                    if os.access(candidate, os.X_OK):
                        test_executable = candidate
                        break
        
        if not test_executable:
            # If not found, try using ctest to run
            print(f"⚠️  Warning: Could not find test executable {self.test_name}, trying ctest...")
            try:
                subprocess.check_call(
                    ["ctest", "-R", self.test_name, "--verbose"],
                    cwd=cmake_dir,
                    env=env
                )
                print(f"✅ Test {self.test_name} passed!")
            except subprocess.CalledProcessError as e:
                print(f"❌ Failed to run test {self.test_name}")
                raise
        else:
            # Run test executable directly
            print(f"🚀 Running test: {test_executable}")
            try:
                subprocess.check_call([test_executable], cwd=os.path.dirname(test_executable), env=env)
                print(f"✅ Test {self.test_name} passed!")
            except subprocess.CalledProcessError as e:
                print(f"❌ Test {self.test_name} failed with exit code {e.returncode}")
                raise

class BuildDistWheel(bdist_wheel):
    user_options = bdist_wheel.user_options + [
        ("device=", None, "target device type (npu or mlu or cuda or ilu or musa)"),
        ("arch=", None, "target arch type (x86 or arm)"),
    ]

    def initialize_options(self) -> None:
        super().initialize_options()
        self.device: Optional[str] = None
        self.arch: Optional[str] = None
        # Cache the original dist name early so finalize_options is idempotent
        # and so name changes are visible to egg_info/metadata generation.
        self._base_dist_name = self.distribution.metadata.name

    def finalize_options(self) -> None:
        # IMPORTANT: mutate distribution name BEFORE super().finalize_options().
        # bdist_wheel finalization may finalize/cache egg_info metadata; if we
        # change the name afterwards, the wheel filename and METADATA can diverge
        # (pip will reject the wheel as "inconsistent name").
        name = self._base_dist_name

        # generate distribution name suffix
        if self.device:
            name += f"_{self.device}"

        torch_version = get_torch_version(self.device)
        if torch_version:
            name += f"_torch{torch_version}"

        if get_cxx_abi():
            name += "_cxx11_abi"
        else:
            name += "_no_cxx11_abi"

        self.distribution.metadata.name = name
        super().finalize_options()

    def run(self) -> None:
        build_ext_cmd = self.get_finalized_command('build_ext')
        build_ext_cmd.device = self.device
        build_ext_cmd.arch = self.arch

        print("🔨 build project...")
        self.run_command('build')

        print("🧪 testing UT...")
        self.run_command('test')

        if self.arch == 'arm':
            ext_path = get_base_dir() + f"/build/lib.linux-aarch64-cpython-{get_python_version()}/"
        else:
            ext_path = get_base_dir() + f"/build/lib.linux-x86_64-cpython-{get_python_version()}/"
        if len(ext_path) == 0:
            print("❌ Build wheel failed, not found path.")
            exit(1)
        tmp_path = os.path.join(ext_path, 'xllm')
        for root, dirs, files in os.walk(tmp_path):
            for item in files:
                path = os.path.join(root, item)
                if '_test' in item and os.path.isfile(path):
                    os.remove(path)
        global BUILD_TEST_FILE
        BUILD_TEST_FILE = False

        self.skip_build = True
        super().run()

class TestUT(Command):
    description = "Run all testing binary."
    user_options = []
    
    # Whitelist: tests that must run sequentially (not in parallel with others)
    # Add test names here if they use fork() or have device initialization conflicts
    # Note: Use test case name patterns (from gtest), not executable names
    SEQUENTIAL_TESTS = [
        'ReduceScatterMultiDeviceTest',
        'DeepEPMultiDeviceTest',
        'AttentionMultiDeviceTest',
        'FusedMoEAll2AllMultiDeviceTest',
    ]

    def initialize_options(self) -> None:
        pass

    def finalize_options(self) -> None:
        pass

    def run_ctest(self, cmake_dir: str) -> int:
        def run_subprocess_with_streaming(
            cmd: list[str],
            error_message: str,
            warn_if_no_tests: bool = False,
        ) -> None:
            """Helper function to run subprocess and stream output"""
            process = subprocess.Popen(
                cmd,
                cwd=cmake_dir,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
            )
            
            if process.stdout is None:
                raise RuntimeError("Failed to capture subprocess stdout for streaming.")

            output_lines: list[str] = []
            for line in iter(process.stdout.readline, ''):
                print(line, end='')
                output_lines.append(line)
            
            return_code: int = process.wait()
            
            # Warn if no tests were found, but don't fail (some backends may not compile certain tests)
            if warn_if_no_tests and return_code == 0:
                output_text: str = ''.join(output_lines)
                if 'No tests were found' in output_text:
                    print(f"No tests matched the pattern (this is OK for some backends).")
                    return
            
            if return_code != 0:
                print(error_message)
                exit(1)
        
        try:
            # Step 1: Run all tests EXCEPT sequential ones in parallel
            if self.SEQUENTIAL_TESTS:
                exclude_pattern = '|'.join(self.SEQUENTIAL_TESTS)
                print("=" * 80)
                print(f"Running tests in parallel (excluding: {', '.join(self.SEQUENTIAL_TESTS)})...")
                print("=" * 80)
                run_subprocess_with_streaming(
                    ['ctest', '--parallel', '8', '--repeat', 'until-pass:5', '-E', exclude_pattern],
                    "Parallel tests failed."
                )
            else:
                print("=" * 80)
                print("Running all tests in parallel...")
                print("=" * 80)
                run_subprocess_with_streaming(
                    ['ctest', '--parallel', '8', '--repeat', 'until-pass:5'],
                    "Parallel tests failed."
                )
            
            # Step 2: Run sequential tests one by one
            for idx, test_name in enumerate(self.SEQUENTIAL_TESTS, start=2):
                print("\n" + "=" * 80)
                print(f"Step {idx}: Running {test_name} sequentially...")
                print("=" * 80)
                # Use pattern matching to include all test cases under the test class
                # e.g., ReduceScatterMultiDeviceTest matches ReduceScatterMultiDeviceTest.BasicTest, etc.
                run_subprocess_with_streaming(
                    ['ctest', '--repeat', 'until-pass:5', '-R', test_name],
                    f"Sequential test {test_name} failed.",
                    warn_if_no_tests=True
                )
            
            print("\n" + "=" * 80)
            print("All tests passed!")
            print("=" * 80)
            return 0
        except subprocess.CalledProcessError as e:
            print(e.stderr)
            exit(1)

    def run(self) -> None:
        self.run_ctest(get_cmake_dir())

class SingleTest(Command):
    """Command to build and run a single test"""
    description = "Build and run a single test target."
    # test_name should match a CMake/CTest target name, for example:
    #   python setup.py test --test-name common_test
    user_options = [
        ("test-name=", None, "name of the test target to build and run (e.g. platform_vmm_test)"),
        ("device=", None, "target device type (npu or mlu or cuda or ilu)"),
        ("arch=", None, "target arch type (x86 or arm)"),
        ("generate-so=", None, "generate so or binary"),
    ]

    def initialize_options(self) -> None:
        self.test_name: Optional[str] = None
        self.device: Optional[str] = None
        self.arch: Optional[str] = None
        self.generate_so: bool = False

    def finalize_options(self) -> None:
        if not self.test_name:
            raise ValueError("--test-name is required for single_test command")

    def run(self) -> None:
        # Create ExtBuildSingleTest instance and set parameters
        build_ext = ExtBuildSingleTest(self.distribution)
        build_ext.initialize_options()
        build_ext.test_name = self.test_name
        build_ext.device = self.device
        build_ext.arch = self.arch
        build_ext.generate_so = self.generate_so
        build_ext.finalize_options()
        
        # Ensure extension modules are set
        if not hasattr(build_ext, 'extensions') or not build_ext.extensions:
            build_ext.extensions = self.distribution.ext_modules
        
        # Run build
        build_ext.run()

def parse_arguments() -> dict[str, Any]:
    parser = argparse.ArgumentParser(
        description='Setup helper for building xllm',
        epilog='Example: python setup.py build',
        usage='%(prog)s [COMMAND] [OPTIONS]'
    )
    
    parser.add_argument(
        'setup_args',
        nargs='*',
        metavar='argparse.REMAINDER',
        help='setup command (build, test, bdist_wheel, etc.)'
    )
    
    parser.add_argument(
        '--device',
        type=str.lower,
        choices=['auto', 'npu', 'mlu', 'cuda', 'ilu', 'musa'],
        default='auto',
        help='Device type: npu, mlu, ilu, cuda or musa (case-insensitive)'
    )
    parser.add_argument(
        '--generate-so',
        type=str.lower,
        choices=['true', 'false', '1', '0', 'yes', 'no', 'y', 'n', 'on', 'off'],
        default='false',
        help='Whether to generate so or binary'
    )
    
    parser.add_argument(
        '--test-name',
        type=str,
        default=None,
        help='Name of the test target to build and run; when omitted, all tests run'
    )

    args = parser.parse_args()
    
    sys.argv = [sys.argv[0]] + args.setup_args
    
    generate_so = args.generate_so.lower() in ('true', '1', 'yes', 'y', 'on')

    return {
        'device': args.device,
        'generate_so': generate_so,
        'test_name': args.test_name,
    }

if __name__ == "__main__":
    config = parse_arguments()

    arch = get_cpu_arch()
    device = config['device']
    if device == 'auto':
        device = get_device_type()
    print(f"🚀 Build xllm with CPU arch: {arch} and target device: {device}")

    pre_build()

    generate_so = config['generate_so']
    test_name = config.get('test_name')

    if "SKIP_TEST" in os.environ:
        BUILD_TEST_FILE = False
    if "SKIP_EXPORT" in os.environ:
        BUILD_EXPORT = False
    
    version = get_version()

    # check and install git pre-commit
    check_and_install_pre_commit()

    test_cmd = SingleTest if test_name else TestUT
    options = {
        'build_ext': {
            'device': device,
            'arch': arch,
            'generate_so': generate_so
        },
        'bdist_wheel': {
            'device': device,
            'arch': arch,
        }
    }
    if test_name:
        options['test'] = {
            'device': device,
            'arch': arch,
            'generate_so': generate_so,
            'test_name': test_name,
        }

    setup(
        name="xllm",
        version=version,
        license="Apache 2.0",
        author="xLLM Team",
        author_email="infer@xllm.ai",
        description="A high-performance inference system for large language models.",
        long_description=read_readme(),
        long_description_content_type="text/markdown",
        url="https://github.com/jd-opensource/xllm",
        project_urls={
            "Homepage": "https://xllm.readthedocs.io/zh-cn/latest/",
            "Documentation": "https://xllm.readthedocs.io/zh-cn/latest/",
        },
        classifiers=[
            "Intended Audience :: Developers",
            "Intended Audience :: Education",
            "Programming Language :: C++",
            "Programming Language :: Python :: 3 :: Only",
            "Programming Language :: Python :: 3.10",
            "Programming Language :: Python :: 3.11",
            "Programming Language :: Python :: 3.12",
            "Operating System :: POSIX",
            "License :: OSI Approved :: Apache Software License",
            "Topic :: Scientific/Engineering",
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
        ],
        ext_modules=[CMakeExtension("xllm", "xllm/")],
        cmdclass={"build_ext": ExtBuild,
                  "test": test_cmd,
                  'bdist_wheel': BuildDistWheel},
        options=options,
        packages=find_namespace_packages(include=["scripts.build_support"]),
        zip_safe=False,
        py_modules=["xllm/launch_xllm", "xllm/__init__",
                    "xllm/pybind/llm", "xllm/pybind/vlm",
                    "xllm/pybind/embedding", "xllm/pybind/util",
                    "xllm/pybind/args", "xllm/pybind/params",
                    "xllm/pybind/errors", "xllm/pybind/mm_utils"],
        python_requires=">=3.10",
    )