""" Copyright (c) 2025-2025 Huawei Technologies Co., Ltd.
sysHAX-adapter is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
    http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
Created: 2025-11-04
Desc: setup.py
"""
import os
import re
import subprocess
import sys
from setuptools import setup, find_packages
from setuptools.command.build_ext import build_ext
from torch.utils.cpp_extension import CppExtension, BuildExtension, include_paths

def get_version():
    """从 sysHAX_adapter/init.py 中读取版本号"""
    init_path = os.path.join(os.path.dirname(__file__), "sysHAX_adapter", "__init__.py")
    
    # 使用正则表达式匹配版本号
    version_pattern = re.compile(r'__version__\s*=\s*["\']([^"\']+)["\']')
    
    try:
        with open(init_path, "r", encoding="utf-8") as f:
            content = f.read()
            match = version_pattern.search(content)
            if match:
                return match.group(1)
    except (FileNotFoundError, IOError):
        pass
    
    # 如果无法读取，返回默认版本
    return "0.1.0"

def get_requirements():
    requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
    if os.path.exists(requirements_path):
        with open(requirements_path, encoding="utf-8") as f:
            return [line.strip() for line in f if line.strip() and not line.startswith("#")]
    return []

def read_readme():
    readme_path = os.path.join(os.path.dirname(__file__), "README.md")
    if os.path.isfile(readme_path):
        with open(readme_path, encoding="utf-8") as f:
            return f.read()
    return ""

# 获取 PyTorch 的包含路径
try:
    import torch
    TORCH_INCLUDE = include_paths()
except ImportError:
    raise RuntimeError("请先安装 torch: pip install torch")

# 获取 Python 包含路径
def get_python_include():
    import sysconfig
    return sysconfig.get_path('include')

# 获取 Python 库路径
def get_python_library():
    import sysconfig
    libdir = sysconfig.get_config_var('LIBDIR')
    if libdir:
        return libdir
    return '/usr/lib64'

# 检测CPU特性
def detect_cpu_features():
    features_line = None
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                if line.startswith("Features"):
                    features_line = line.strip()
                    break
    except Exception:
        pass
    
    if not features_line:
        print("⚠️ 无法读取 CPU 特性")
        return {"fp16": False, "dotprod": False, "i8mm": False}
    
    features = set(features_line.split()[1:])
    
    has_fp16 = "fphp" in features and "asimdhp" in features
    has_dotprod = "asimddp" in features
    has_i8mm = "i8mm" in features
    
    print(f"✅ CPU Features -> fp16: {has_fp16}, dotprod (asimddp): {has_dotprod}, i8mm: {has_i8mm}")
    return {"fp16": has_fp16, "dotprod": has_dotprod, "i8mm": has_i8mm}

# 获取编译参数
def get_extra_compile_args_and_macros():
    import platform
    args = []
    macros = []
    
    machine = platform.machine()
    if machine not in ['aarch64', 'arm64']:
        print("📎 非 ARM64 架构，跳过优化")
        return args, macros
    
    cpu_feat = detect_cpu_features()
    
    # 构建 -march 标志
    march_parts = ["armv8.2-a"]
    
    # 必须添加 dotprod 才能使用 vdotq_s32
    if cpu_feat["dotprod"]:
        march_parts.append("dotprod")
        macros.append(("ARM_DOTPROD_SUPPORT", "1"))
    else:
        # 即使CPU不支持，也要启用dotprod编译选项，让代码能编译通过
        # 运行时可以通过CPU特性检测来决定是否使用dotprod指令
        march_parts.append("dotprod")
        macros.append(("ARM_DOTPROD_SUPPORT", "1"))
    
    if cpu_feat["fp16"]:
        march_parts.append("fp16")
        macros.append(("ARM_FP16_SUPPORT", "1"))
    
    if cpu_feat["i8mm"]:
        march_parts.append("i8mm")
        macros.append(("ARM_I8MM_SUPPORT", "1"))
    
    march_flag = "-march=" + "+".join(march_parts)
    args.append(march_flag)
    print(f"🔧 使用编译选项: {march_flag}")
    
    # 添加必要的编译标志
    args.extend([
        '-std=c++17',
        '-D_GLIBCXX_USE_CXX11_ABI=1',
        '-fPIC',
        '-O3',
        '-fopenmp',
        '-Wall',
        '-Wextra',
        '-fno-omit-frame-pointer',
        # 确保内联优化生效
        '-finline-functions',
        '-finline-small-functions',
        '-findirect-inlining',
        '-finline-atomics',
    ])
    
    # 如果是GCC，添加特定优化选项
    try:
        compiler_version = subprocess.check_output(['gcc', '-dumpversion']).decode().strip()
        if compiler_version >= '10':
            # GCC 10+ 对ARM有更好的支持
            args.extend([
                '-ftree-vectorize',
                '-fopt-info-vec-optimized',
            ])
    except:
        pass
    
    return args, macros

# 自定义构建扩展类
class CustomBuildExtension(BuildExtension):
    def build_extensions(self):
        # 确保编译器标志应用到所有文件
        extra_compile_args, define_macros = get_extra_compile_args_and_macros()
        
        for ext in self.extensions:
            # 为每个扩展设置编译参数
            ext.extra_compile_args = extra_compile_args
            ext.define_macros = define_macros
            
            # 添加必要的库路径
            torch_lib = os.path.join(os.path.dirname(torch.__file__), 'lib')
            if os.path.exists(torch_lib):
                ext.library_dirs.append(torch_lib)
            
            # 确保链接正确的库
            ext.libraries.extend(['torch', 'torch_cpu', 'c10', 'gomp'])
        
        # 调用父类方法
        super().build_extensions()

# 获取动态编译参数
extra_compile_args, define_macros = get_extra_compile_args_and_macros()

# 定义扩展模块
ext_modules = [
    CppExtension(
        name="cpu_inference",
        sources=[
            "csrc/cpu/cpu_inference.cpp",
            "csrc/cpu/cpu_bindings.cpp",
            "csrc/cpu/cpu_utils.cpp",
            "csrc/cpu/cpu_inference_manager.cpp",
            "csrc/cpu/model_weight_base.cpp",
            "csrc/cpu/qwen3_moe.cpp",
            "csrc/cpu/tensor.cpp",
            "csrc/cpu/quantization/quantization_fp16.cpp",
            "csrc/cpu/quantization/quantization_q4_0.cpp",
            "csrc/cpu/quantization/quantization_q8_0.cpp",
            "csrc/cpu/quantization/quantization_q8align.cpp",
            "csrc/cpu/matmul/matmul_fp16.cpp",
            "csrc/cpu/matmul/matmul_q8.cpp",
            "csrc/cpu/matmul/matmul_q8align.cpp",
            "csrc/cpu/matmul/matmul_q4q8.cpp",
            "csrc/cpu/merge_silu.cpp",
            "csrc/cpu/memory_manager.cpp"
        ],
        include_dirs=[
            ".",
            "csrc",
            "csrc/cpu",
            *TORCH_INCLUDE,
            get_python_include(),
        ],
        library_dirs=[
            get_python_library(),
        ],
        libraries=['python3', 'numa'],
        extra_compile_args=extra_compile_args,
        define_macros=define_macros,
        language='c++',
        extra_link_args=[
            '-fopenmp',
            '-lgomp',
            '-Wl,--no-as-needed',  # 确保所有库都被链接
        ],
    ),
]

# 获取版本号
version = get_version()

setup(
    name="sysHAX-adapter",
    version=version,
    author="openEuler sysHAX team",
    license="Mulan PSL v2",
    description="sysHAX-adapter - Inference framework and inference card adapter",
    long_description=read_readme(),
    long_description_content_type="text/markdown",
    url="https://gitee.com/openeuler/sysHAX-adapter",
    project_urls={
        "Homepage": "https://gitee.com/openeuler/sysHAX-adapter",
    },
    entry_points={
        "console_scripts": ["sysHAX-adapter=sysHAX_adapter.entrypoints:main"]
    },
    classifiers=[
        "Intended Audience :: Developers",
        "License :: OSI Approved :: Mulan Permissive Software License v2",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    packages=find_packages(include=['sysHAX_adapter', 'sysHAX_adapter.*']),
    package_dir={'': '.'},
    python_requires=">=3.9",
    install_requires=get_requirements(),
    include_package_data=True,
    zip_safe=False,
    ext_modules=ext_modules,
    cmdclass={
        "build_ext": CustomBuildExtension,
    },
)