import os
import re
from setuptools import setup
from setuptools.extension import Extension
from Cython.Build import cythonize

def build_current_dir():
    current_dir = os.path.dirname(os.path.abspath(__file__))

    # 要排除的 .py 文件（编译脚本、__init__.py 等）
    exclude_files = {"build_all_to_so.py", "__init__.py"}

    # 获取当前目录下所有可编译的 .py 文件
    py_files = [
        f for f in os.listdir(current_dir)
        if f.endswith(".py") and f not in exclude_files
    ]

    if not py_files:
        print("❗ 当前目录下没有需要编译的 .py 文件。")
        return

    # 创建 Extension 对象列表
    extensions = [
        Extension(
            name=os.path.splitext(f)[0],
            sources=[os.path.join(current_dir, f)]
        )
        for f in py_files
    ]

    try:
        setup(
            script_args=["build_ext", "--inplace"],
            ext_modules=cythonize(
                extensions,
                compiler_directives={"language_level": 3},
            ),
        )
        print("✅ 编译完成")
    except Exception as e:
        print("❌ 编译出错:", e)

    # 可选：删除生成的中间 .c 文件
    for f in py_files:
        c_file = os.path.join(current_dir, f.replace(".py", ".c"))
        if os.path.exists(c_file):
            os.remove(c_file)

def rename_so_files():
    current_dir = os.path.dirname(os.path.abspath(__file__))
    pattern = re.compile(r"^(?P<name>.+?)\.cpython-[\d]+-[\w-]+\.so$")

    for f in os.listdir(current_dir):
        if f.endswith(".so"):
            m = pattern.match(f)
            if m:
                new_name = m.group("name") + ".so"
                old_path = os.path.join(current_dir, f)
                new_path = os.path.join(current_dir, new_name)
                print(f"Rename: {f} -> {new_name}")
                os.rename(old_path, new_path)

if __name__ == "__main__":
    build_current_dir()
    rename_so_files()