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 = "repositoryZN"
COMPILE_DIRS = ["utils"]           # 编译 utils 目录和根目录下的 .py
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. 编译所有需要编译的目录（包括根目录 "."）
        for compile_dir in COMPILE_DIRS:
            if compile_dir == ".":
                target_dir = src_root
            else:
                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. 手动复制：对于 utils 目录作为顶级包，对于根目录下的文件直接复制到 build/lib 根目录
        build_lib = self.build_lib
        if not build_lib:
            build_lib = os.path.join(src_root, "build", "lib")

        # 处理 utils 目录（作为包）
        utils_src = os.path.join(src_root, "utils")
        if os.path.isdir(utils_src):
            utils_dst = os.path.join(build_lib, "utils")
            self.mkpath(utils_dst)
            for root, _, files in os.walk(utils_src):
                rel_path = os.path.relpath(root, utils_src)
                target_subdir = os.path.join(utils_dst, rel_path) if rel_path != '.' else utils_dst
                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}")

        # 处理根目录下的文件（作为顶级模块，直接放在 build/lib 根目录）
        for file in os.listdir(src_root):
            if file.endswith(".pyd") or file.endswith(".so"):
                src_file = os.path.join(src_root, file)
                dst_file = os.path.join(build_lib, file)
                self.copy_file(src_file, dst_file)
                log(f"复制根模块 {src_file} -> {dst_file}")
            elif file == "__init__.py":
                # 如果根目录有 __init__.py，也复制（让根目录本身成为一个包）
                src_file = os.path.join(src_root, file)
                dst_file = os.path.join(build_lib, file)
                self.copy_file(src_file, dst_file)
                log(f"复制根 __init__.py {src_file} -> {dst_file}")

        # 3. 清理源码目录中的 .pyd/.so
        if CLEAN_COMPILED_AFTER_BUILD:
            for compile_dir in COMPILE_DIRS:
                if compile_dir == ".":
                    target_dir = src_root
                else:
                    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()

    # 动态确定要安装的包和模块
    packages = []
    if os.path.isdir("utils"):
        packages.append("utils")
    # 根目录下的顶级模块（如果存在这些 .py 文件，它们会被编译成 .pyd 并安装）
    py_modules = []
    for mod in ["jsonUtil", "logUtil", "csvFileUtil"]:
        if os.path.exists(f"{mod}.py"):
            py_modules.append(mod)

    setup(
        name=PACKAGE_NAME,
        version="1.0.202508190930",
        description="Repository ZN Package (compiled, flat structure)",
        author="Xie Zhou Yang",
        packages=packages,
        py_modules=py_modules,
        package_dir={"": "."},
        cmdclass={
            "build_py": CustomBuildPy,
            "bdist_wheel": CustomBdistWheel,
        },
        install_requires=INSTALL_REQUIRES,
        zip_safe=False,
    )