import os
import shutil
from distutils.core import setup
from Cython.Build import cythonize

"""运行方式: 
python setup.py build_ext --inplace
"""

# 遍历搜索目录下所有".py"文件并删除"__pycache__"文件夹
def path_search(root_path, pylist, exclude_list, suffix=".py"):
    # 删除"__pycache__"文件夹下".pyc"文件,防止被反编译
    pycache_path = os.path.join(root_path, "__pycache__")
    shutil.rmtree(pycache_path) if os.path.exists(pycache_path) else None
    # 获取当前目录下的所有文件集合
    file_list = os.listdir(root_path)
    # 遍历搜索目录下所有".py"文件
    for file in file_list:
        # 路径拼接
        path = os.path.join(root_path, file)
        # 若path是目录则,向下搜索
        if os.path.isdir(path):
            pylist = path_search(path, pylist, exclude_list)
        # 判断文件名是否以".py"结尾,并且不是排除文件
        elif file.endswith(suffix) and file not in exclude_list:
            pylist.append(path)
    return pylist

# 生成cython加密文件
def setup_func(pylist, exclude_list):
    # 注意这里推荐使用相对路径,编译出的so文件在引用其他模块时可能会出现路径问题
    setup(ext_modules=cythonize(pylist, exclude=exclude_list))

# 清理setup生成的临时文件
def clean_setup_file(root_path, pylist):
    # 删除"build"文件夹
    shutil.rmtree("build")
    # 获取主目录下的所有".so"或".pyd"文件集合
    file_list = os.listdir(root_path)
    cylist = [f for f in file_list if f.endswith((".so", ".pyd"))]
    # 删除非排除".py"和".c"文件
    for pypath in pylist:
        os.remove(pypath)
        os.remove(pypath[:-3] + ".c")
        # ".py"文件路径分割
        split_pypath = os.path.split(pypath)
        # 移动".so"或".pyd"文件到相应文件夹中
        for cyfile in cylist:
            if cyfile.startswith(split_pypath[1][:-3]) is True:
                # 判断".so"或".pyd"文件路径是否与根路径相同,若不同则移动文件
                if os.path.samefile(split_pypath[0], root_path) is False:
                    shutil.move(cyfile, split_pypath[0])
                break
    # 删除"setup.py"文件
    os.remove("setup.py")


if __name__ == "__main__":
    # 获取当前路径
    root_path = os.getcwd()
    # 预先复制原文件
    split_path = os.path.split(root_path)
    copy_path = os.path.join(split_path[0], split_path[1] + "_copy")
    if os.path.exists(copy_path) is False:
        shutil.copytree(root_path, copy_path)
    
    # 非加密".py"文件列表
    exclude_list=["setup.py", "__init__.py", "config.py", "guc_conf.py", 
                  "retrieval_monitor.py", "restart_server.py"]
    # 获取当前目录下的所有"app.py"文件加入exclude_list列表
    os.chdir(copy_path)
    file_list = os.listdir(copy_path)
    exclude_list.extend([f for f in file_list if f.endswith("app.py")])
    
    # 需要setup加密的".py"文件列表
    pylist = []
    # 遍历搜索目录下所有py文件并返回列表
    pylist = path_search(copy_path, pylist, exclude_list)
    # 生成cython加密文件
    setup_func(pylist, exclude_list)
    # 清理setup生成的临时文件
    clean_setup_file(copy_path, pylist)