import os
import sys
import re
import shutil
import subprocess
from setuptools import setup, find_packages
from setuptools.command.build_py import build_py
from setuptools.command.bdist_wheel import bdist_wheel

# ========== 配置 ==========
PACKAGE_NAME = "dataAnalysisBusiness"
COMPILE_DIRS = ["algorithm"]          # 需要编译的子目录（将作为顶层包）
# INSTALL_REQUIRES = ["dataAnalysisBehavior"]   # 根据实际依赖填写
INSTALL_REQUIRES = []
CLEAN_COMPILED_AFTER_BUILD = True
# =========================

def log(msg):
    print(f"[SETUP] {msg}")

def compile_with_nuitka(py_files):
    for py_file in py_files:
        if os.path.basename(py_file) == "__init__.py":
            continue
        log(f"编译: {py_file}")

        output_dir = os.path.dirname(py_file)
        cmd = [
            sys.executable, "-m", "nuitka",
            "--module",
            "--no-pyi-file",
            "--remove-output",
            f"--output-dir={output_dir}",
            py_file
        ]
        if sys.platform == "win32":
            cmd.append("--mingw64")
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            log(f"编译失败: {py_file}\n{result.stderr}")
            sys.exit(1)

        base_name = os.path.splitext(os.path.basename(py_file))[0]
        found = False
        for ext in [".pyd", ".so"]:
            pattern = re.escape(base_name) + r".*" + re.escape(ext)
            for f in os.listdir(output_dir):
                if re.match(pattern, f):
                    src = os.path.join(output_dir, f)
                    dst = os.path.join(output_dir, base_name + ext)
                    shutil.move(src, dst)
                    log(f"重命名: {src} -> {dst}")
                    found = True
                    break
            if found:
                break
        if not found:
            log(f"错误：未找到编译产物 {py_file}")
            sys.exit(1)

def clean_build_dir():
    dirs_to_clean = ["build", f"{PACKAGE_NAME}.egg-info"]
    for d in dirs_to_clean:
        if os.path.exists(d):
            shutil.rmtree(d)
            log(f"已删除: {d}")

class CustomBuildPy(build_py):
    def run(self):
        src_root = os.getcwd()
        log(f"源码根目录: {src_root}")

        # 1. 编译生成 .pyd/.so
        for compile_dir in COMPILE_DIRS:
            target_dir = os.path.join(src_root, compile_dir)
            if not os.path.isdir(target_dir):
                log(f"警告: 目录不存在 {target_dir}")
                continue
            all_py = []
            for root, _, files in os.walk(target_dir):
                for file in files:
                    if file.endswith(".py"):
                        all_py.append(os.path.join(root, file))
            if all_py:
                compile_with_nuitka(all_py)
            else:
                log(f"在 {target_dir} 中没有找到 .py 文件")

        # 2. 手动构建：复制到 build/lib/{compile_dir}（直接作为顶级包）
        build_lib = self.build_lib
        if not build_lib:
            build_lib = os.path.join(src_root, "build", "lib")

        for compile_dir in COMPILE_DIRS:
            src_dir = os.path.join(src_root, compile_dir)
            if not os.path.isdir(src_dir):
                continue
            dst_dir = os.path.join(build_lib, compile_dir)   # 顶级包
            self.mkpath(dst_dir)
            for root, _, files in os.walk(src_dir):
                rel_path = os.path.relpath(root, src_dir)
                target_subdir = os.path.join(dst_dir, rel_path) if rel_path != '.' else dst_dir
                self.mkpath(target_subdir)
                for file in files:
                    src_file = os.path.join(root, file)
                    if file == "__init__.py":
                        dst_file = os.path.join(target_subdir, file)
                        self.copy_file(src_file, dst_file)
                        log(f"复制 {src_file} -> {dst_file}")
                    elif file.endswith(".pyd") or file.endswith(".so"):
                        dst_file = os.path.join(target_subdir, file)
                        self.copy_file(src_file, dst_file)
                        log(f"复制 {src_file} -> {dst_file}")

        # 3. 清理源码目录中的 .pyd/.so
        if CLEAN_COMPILED_AFTER_BUILD:
            for compile_dir in COMPILE_DIRS:
                target_dir = os.path.join(src_root, compile_dir)
                if not os.path.isdir(target_dir):
                    continue
                for root, _, files in os.walk(target_dir):
                    for file in files:
                        if file.endswith(".pyd") or file.endswith(".so"):
                            os.remove(os.path.join(root, file))
                            log(f"清理源码编译产物: {os.path.join(root, file)}")

class CustomBdistWheel(bdist_wheel):
    def run(self):
        self.run_command("build_py")
        self.plat_name = self.get_platform()
        bdist_wheel.run(self)

    def get_platform(self):
        if sys.platform == "win32":
            return "win_amd64"
        elif sys.platform == "linux":
            return "linux_x86_64"
        else:
            return super().get_platform()

if __name__ == "__main__":
    try:
        subprocess.run([sys.executable, "-m", "nuitka", "--version"], capture_output=True, check=True)
    except subprocess.CalledProcessError:
        log("请先安装 Nuitka: pip install nuitka", "ERROR")
        sys.exit(1)

    clean_build_dir()

    setup(
        name=PACKAGE_NAME,
        version="1.2.202508190930",
        description="Data Analysis Business Package (compiled, flat structure)",
        author="Xie Zhou Yang",
        packages=["algorithm"],                # 顶层包
        package_dir={"": "."},
        cmdclass={
            "build_py": CustomBuildPy,
            "bdist_wheel": CustomBdistWheel,
        },
        install_requires=INSTALL_REQUIRES,
        zip_safe=False,
    )