# setup.py
from setuptools import setup, Extension
from Cython.Build import cythonize
from Cython.Distutils import build_ext
import os
import sys
from pathlib import Path
import re

# 获取当前目录
current_dir = Path(__file__).parent
COMPILED_DIR = current_dir / "build" / "compiled"

# 定义要排除的文件
EXCLUDE_FILES = {
    "app.py",  # 不编译主应用文件
    "setup.py",  # 不编译本文件
    "cythonize.py",  # 不编译辅助脚本
    "check_deploy.py",  # 不编译部署检查
    "terminal_demo.py",  # 不编译终端演示文件
}

# 定义要排除的目录
EXCLUDE_DIRS = {
    "build",
    "dist",
    "__pycache__",
    ".git",
    ".vscode",
    ".idea",
    "venv",
    "env",
    "node_modules",
    "chat_logs",
    "logs",
    "tests",
    "test",
    "cython_build",  # Cython临时目录
}


def should_exclude_dir(dir_name):
    """判断是否应该排除目录"""
    return dir_name in EXCLUDE_DIRS or dir_name.startswith(".")


def find_py_files(base_dir):
    """查找所有需要编译的Python文件"""
    py_files = []

    for root, dirs, files in os.walk(base_dir):
        # 排除不需要的目录
        dirs[:] = [d for d in dirs if not should_exclude_dir(d)]

        for file in files:
            if file.endswith(".py") and file not in EXCLUDE_FILES:
                full_path = os.path.join(root, file)

                # 排除build目录下的文件
                if "build" in full_path.split(os.sep):
                    continue

                py_files.append(full_path)

    return py_files


def create_extension_modules():
    """创建Cython扩展模块列表"""
    extensions = []
    py_files = find_py_files(current_dir)

    for py_file in py_files:
        # 转换为相对路径
        rel_path = os.path.relpath(py_file, current_dir)

        # 生成模块名（去掉.py）
        module_name = rel_path[:-3].replace(os.sep, ".")

        # 创建Extension
        extension = Extension(
            module_name,
            [py_file],
            extra_compile_args=[
                "/O2" if sys.platform == "win32" else "-O3",
                "/std:c++17" if sys.platform == "win32" else "-std=c++17",
            ],
            language="c",
        )

        extensions.append(extension)

    return extensions


def main():
    """主编译函数"""
    print(f"编译目标目录: {COMPILED_DIR}")

    # 确保编译目录存在
    COMPILED_DIR.mkdir(parents=True, exist_ok=True)

    # 创建扩展模块
    extensions = create_extension_modules()

    if not extensions:
        print("❌ 未找到需要编译的Python文件")
        return

    print(f"找到 {len(extensions)} 个Python文件需要编译:")
    for i, ext in enumerate(extensions[:10], 1):  # 只显示前10个
        print(f"  {i:2}. {ext.name}")
    if len(extensions) > 10:
        print(f"  ... 和 {len(extensions) - 10} 个其他文件")

    # 配置编译选项
    compiler_directives = {
        "language_level": 3,
        "boundscheck": False,  # 关闭边界检查
        "wraparound": False,  # 关闭环绕检查
        "initializedcheck": False,  # 关闭初始化检查
        "nonecheck": False,  # 关闭None检查
        "overflowcheck": False,  # 关闭溢出检查
        "cdivision": True,  # 使用C除法
        "infer_types": True,  # 类型推断
        "optimize.use_switch": True,  # 优化switch语句
    }

    # 编译
    cythonized = cythonize(
        extensions,
        compiler_directives=compiler_directives,
        nthreads=0,  # 0表示自动使用所有核心
        build_dir="./cython_build",  # 临时构建目录
    )

    # 修复：使用自定义BuildExt类
    class CustomBuildExt(build_ext):
        def initialize_options(self):
            super().initialize_options()
            # 设置输出目录
            self.build_lib = str(COMPILED_DIR)

        def get_ext_fullpath(self, ext_name):
            """获取扩展的完整路径，确保输出到指定目录"""
            # 首先调用父类方法
            filename = self.get_ext_filename(ext_name)

            # 确保输出到编译目录
            return os.path.join(self.build_lib, filename)

    # 修复：使用setuptools的setup，但通过cmdclass和options控制
    setup(
        name="LongjoeAgent",
        ext_modules=cythonized,
        cmdclass={"build_ext": CustomBuildExt},
        options={
            "build_ext": {
                "build_lib": str(COMPILED_DIR),
                "inplace": False,  # 不在原地构建
            }
        },
        script_args=["build_ext"],  # 只构建扩展
    )


if __name__ == "__main__":
    main()