import logging
import os
import shutil
import subprocess
import sys
from sysconfig import get_paths

from setuptools import setup, find_packages, Extension, Command
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
from setuptools.command.develop import develop
from setuptools.command.install import install


def req_file(filename):
    with open(filename, encoding='utf-8') as f:
        content = f.readlines()
    return [x.strip() for x in content]


install_requires = req_file("requirements.txt")

ROOT_DIR = os.path.dirname(__file__)


def get_value_from_lines(lines: list[str], key: str) -> str:
    for line in lines:
        line = " ".join(line.split())
        if key in line:
            return line.split(":")[-1].strip()
    return ""


def get_chip_type() -> str:
    npu_smi = shutil.which("npu-smi")
    if npu_smi is None:
        logging.warning(
            "npu-smi command not found, if this is an npu environment, "
            "please check if npu driver is installed correctly."
        )
        return ""
    try:
        npu_info_lines = subprocess.check_output([npu_smi, "info", "-l"]).decode().strip().split("\n")
        npu_id = int(get_value_from_lines(npu_info_lines, "NPU ID"))

        board_info_lines = (
            subprocess.check_output([npu_smi, "info", "-t", "board", "-i", str(npu_id)]).decode().strip().split("\n")
        )

        chip_name = get_value_from_lines(board_info_lines, "Chip Name")

        if not chip_name:
            chip_info_lines = (
                subprocess.check_output([npu_smi, "info", "-t", "board", "-i", str(npu_id), "-c", "0"])
                .decode()
                .strip()
                .split("\n")
            )
        else:
            chip_info_lines = board_info_lines

        chip_name = get_value_from_lines(chip_info_lines, "Chip Name")
        chip_type = get_value_from_lines(chip_info_lines, "Chip Type")
        npu_name = get_value_from_lines(chip_info_lines, "NPU Name")

        if "910" in chip_name:
            if chip_type:
                # arch32 case
                assert not npu_name
                return (chip_type + chip_name).lower()
            else:
                # arch32 case
                assert npu_name
                return (chip_name + "_" + npu_name).lower()
        elif "950" in chip_name:
            # arch35 case
            assert npu_name
            return (chip_name + "_" + npu_name).lower()
        else:
            raise ValueError(f"Unable to recognize chip name: {chip_name}, please manually set env SOC_VERSION")
    except subprocess.CalledProcessError as e:
        logging.warning("npu-smi call failed: %s. ACLNN custom operator compilation will be skipped.", e)
        return ""
    except FileNotFoundError:
        logging.warning(
            "npu-smi command not found, if this is an npu environment, "
            "please check if npu driver is installed correctly."
        )
        return ""


_env_soc_version = os.environ.get("SOC_VERSION", "")
if not _env_soc_version:
    _env_soc_version = get_chip_type()
    if _env_soc_version:
        logging.info("Auto-detected SOC_VERSION: %s", _env_soc_version)
    else:
        logging.warning(
            "Could not determine chip type automatically. "
            "ACLNN custom operator compilation will be skipped.\n"
            "To enable aclnn compilation, set the 'SOC_VERSION' environment "
            "variable to specify the target chip, for example:\n"
            '  - arch32: export SOC_VERSION="ascend910b1"\n'
            '  - arch32: export SOC_VERSION="ascend910_9391"\n'
            '  - arch35: export SOC_VERSION="<value starting with ascend950>"\n'
        )


def _find_cmake_binary():
    cmake_from_path = shutil.which("cmake")
    if cmake_from_path and os.path.isfile(cmake_from_path):
        try:
            with open(cmake_from_path, "rb") as f:
                head = f.read(128)
            if not head.startswith(b"#!/") or b"python" not in head.lower():
                return cmake_from_path
        except OSError:
            pass
    for candidate in ["/usr/bin/cmake", "/usr/local/bin/cmake"]:
        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    if cmake_from_path:
        return cmake_from_path
    return "cmake"


class CMakeExtension(Extension):
    def __init__(self, name: str, cmake_lists_dir: str = ".", **kwargs) -> None:
        super().__init__(name, sources=[], py_limited_api=False, **kwargs)
        self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)


class cmake_build_ext(build_ext):
    did_config: dict[str, bool] = {}

    def compute_num_jobs(self):
        num_jobs = os.environ.get("MAX_JOBS")
        if num_jobs is not None:
            num_jobs = int(num_jobs)
        else:
            try:
                num_jobs = len(os.sched_getaffinity(0))
            except AttributeError:
                num_jobs = os.cpu_count()
        num_jobs = max(1, num_jobs)
        return num_jobs

    def configure(self, ext: CMakeExtension) -> None:
        build_temp = self.build_temp
        os.makedirs(build_temp, exist_ok=True)
        source_dir = os.path.abspath(ROOT_DIR)
        python_executable = sys.executable

        cmake_args = [_find_cmake_binary()]
        cmake_args += ["-DCMAKE_BUILD_TYPE=Release"]
        cmake_args += ["-DCMAKE_EXPORT_COMPILE_COMMANDS=1"]

        asc_home = os.environ.get("ASCEND_HOME_PATH")
        if asc_home is None:
            logging.warning(
                "No ASCEND_HOME_PATH found in environment. "
                "Using default path. Set ASCEND_HOME_PATH if you have a "
                "custom CANN installation."
            )
            asc_home = "/usr/local/Ascend/ascend-toolkit/latest"
        logging.info("ASCEND_HOME_PATH: %s", asc_home)
        cmake_args += [f"-DASCEND_HOME_PATH={asc_home}"]
        cmake_args += [f"-DPYTHON_EXECUTABLE={python_executable}"]
        cmake_args += [f"-DPYTHON_INCLUDE_PATH={get_paths()['include']}"]

        # Detect pybind11 CMake directory.
        # If pybind11 is installed via pip, this will return its cmake path.
        try:
            pybind11_cmake_path = (
                subprocess.check_output(
                    [python_executable, "-m", "pybind11", "--cmakedir"],
                    stderr=subprocess.PIPE,
                )
                .decode()
                .strip()
            )
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"Failed to locate pybind11 CMake directory via "
                f"'{python_executable} -m pybind11 --cmakedir'. "
                f"Install pybind11 in this Python environment:\n"
                f"  {python_executable} -m pip install pybind11"
            ) from e
        cmake_args += [f"-DCMAKE_PREFIX_PATH={pybind11_cmake_path}"]

        install_path = os.path.join(ROOT_DIR, self.build_lib)
        if isinstance(self.distribution.get_command_obj("develop"), develop):
            install_path = os.path.join(ROOT_DIR, "mindspeed_ops")
        cmake_args += [f"-DCMAKE_INSTALL_PREFIX={install_path}"]

        # Detect torch_npu installation path.
        # Prefer import-then-inspect over pip-show: pip-show may fail in
        # pip's build-isolation environment even when torch-npu IS installed
        # there (via pyproject.toml requires).
        try:
            torch_npu_path = (
                subprocess.check_output(
                    [python_executable, "-c", "import torch_npu; print(torch_npu.__file__)"],
                    stderr=subprocess.PIPE,
                )
                .decode()
                .strip()
            )
            # __file__ → .../torch_npu/__init__.py, take the dirname
            torch_npu_path = os.path.dirname(torch_npu_path)
        except subprocess.CalledProcessError:
            torch_npu_path = ""
            logging.warning(
                "torch_npu not found via import. "
                "If torch_npu is installed at a non-standard location, "
                "set TORCH_NPU_PATH manually."
            )
        if torch_npu_path:
            cmake_args += [f"-DTORCH_NPU_PATH={torch_npu_path}"]

        cmake_args += [source_dir]

        logging.info("cmake config command: %s", cmake_args)
        try:
            subprocess.check_call(cmake_args, cwd=self.build_temp)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"CMake configuration failed: {e}")

    def build_extensions(self) -> None:
        try:
            subprocess.check_output([_find_cmake_binary(), "--version"])
        except OSError as e:
            raise RuntimeError(f"Cannot find CMake executable: {e}")

        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

        targets = []

        def target_name(s: str) -> str:
            return s.removeprefix("mindspeed_ops.")

        for ext in self.extensions:
            self.configure(ext)
            targets.append(target_name(ext.name))

        num_jobs = self.compute_num_jobs()
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
        try:
            subprocess.check_call([_find_cmake_binary(), *build_args], cwd=self.build_temp)
        except OSError as e:
            raise RuntimeError(f"Build library failed: {e}")

        install_args = [
            _find_cmake_binary(),
            "--install",
            ".",
        ]
        try:
            subprocess.check_call(install_args, cwd=self.build_temp)
        except OSError as e:
            raise RuntimeError(f"Install library failed: {e}")

        # Copy _cann_ops_custom from source tree to build output.
        # The directory is populated at install time by build_aclnn.sh
        # and must be present alongside the .so for
        # ASCEND_CUSTOM_OPP_PATH autoconfig to work.
        src_cann = os.path.join(ROOT_DIR, "mindspeed_ops", "_cann_ops_custom")
        dst_cann = os.path.join(self.build_lib, "mindspeed_ops", "_cann_ops_custom")
        if os.path.isdir(src_cann):
            if os.path.exists(dst_cann):
                shutil.rmtree(dst_cann)
            shutil.copytree(src_cann, dst_cann)
            print(f"Copy: {src_cann} -> {dst_cann}")

        if isinstance(self.distribution.get_command_obj("develop"), develop):
            for root, _, files in os.walk(self.build_temp):
                for file in files:
                    if file.endswith(".so"):
                        src_path = os.path.join(root, file)
                        dst_path = os.path.join(self.build_lib, "mindspeed_ops", file)
                        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
                        shutil.copy(src_path, dst_path)
                        print(f"Copy: {src_path} -> {dst_path}")

    def run(self):
        if _env_soc_version:
            self.run_command("build_aclnn")
        super().run()


class build_and_install_aclnn(Command):
    description = "Build and install AclNN by running build_aclnn.sh"
    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        try:
            print("Running bash build_aclnn.sh ...")
            subprocess.check_call(
                [shutil.which("bash") or "bash", "mindspeed_ops/csrc/build_aclnn.sh", ROOT_DIR, _env_soc_version]
            )
            print("build_aclnn.sh executed successfully!")
        except subprocess.CalledProcessError as e:
            print(f"Error running build_aclnn.sh: {e}")
            raise SystemExit(e.returncode)


class custom_develop(develop):
    def run(self):
        self.run_command("build_ext")
        super().run()


class custom_install(install):
    def run(self):
        self.run_command("build_ext")
        install.run(self)


ext_modules = []
if _env_soc_version:
    ext_modules = [CMakeExtension(name="mindspeed_ops.mindspeed_ops_C")]

cmdclass = {
    "develop": custom_develop,
    "build_py": build_py,
    "build_aclnn": build_and_install_aclnn,
    "build_ext": cmake_build_ext,
    "install": custom_install,
}

setup(
    name="mindspeed_ops",
    version="0.1.0",
    description="Ascend-optimized custom operators for training tasks",
    packages=find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires='>=3.10',
    install_requires=install_requires,
    ext_modules=ext_modules,
    cmdclass=cmdclass,
)