import os
import subprocess

def install_cuda_toolkit():
    try:
        print("正在安装 CUDA...")
        CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
        CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)

        # 下载 CUDA 安装包并显示进度
        print(f"下载 CUDA 安装包: {CUDA_TOOLKIT_URL}")
        subprocess.run(["wget", "--progress=bar:force", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE], check=True)

        # 给安装包添加执行权限并显示输出
        print(f"为安装包添加执行权限: {CUDA_TOOLKIT_FILE}")
        subprocess.run(["chmod", "+x", CUDA_TOOLKIT_FILE], check=True)

        # 安装 CUDA 12.1，并显示安装过程的输出
        print("正在安装 CUDA...")
        subprocess.run([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        # 设置环境变量
        os.environ["CUDA_HOME"] = "/usr/local/cuda"
        os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
        os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
            os.environ["CUDA_HOME"],
            "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
        )
        
        # 设置 PyTorch 使用的 CUDA 架构
        os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;9.0"  # 可以根据需要调整
        print("CUDA 12.1 安装并配置完成！")
    
    except subprocess.CalledProcessError as e:
        print(f"安装 CUDA 失败: {e}")
        print(f"错误输出: {e.stderr}")

def install_custom_rasterizer():
    try:
        print("正在编译并安装 custom_rasterizer...")
        # 假设 custom_rasterizer 代码目录在 /path/to/your/project/texgen/custom_rasterizer
        custom_rasterizer_dir = "/home/user/app/hy3dgen/texgen/custom_rasterizer"  # 替换为实际路径
        
        # 进入 custom_rasterizer 目录
        os.chdir(custom_rasterizer_dir)
        
        # 编译并生成 .whl 文件，并显示输出
        print("正在编译 custom_rasterizer...")
        subprocess.run(["python3", "setup.py", "bdist_wheel"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        # 获取生成的 .whl 文件路径
        whl_file_path = os.path.join(custom_rasterizer_dir, "dist", "custom_rasterizer-<version>-cp310-cp310-linux_x86_64.whl")
        
        print(f"custom_rasterizer 的 .whl 文件已生成：{whl_file_path}")
        return whl_file_path
    
    except subprocess.CalledProcessError as e:
        print(f"安装 custom_rasterizer 失败: {e}")
        print(f"错误输出: {e.stderr}")

# 主程序
if __name__ == "__main__":
    try:
        print("正在安装 CUDA")
        # 安装 CUDA 12.1
        install_cuda_toolkit()
        
        print("正在安装 custom_rasterizer")
        # 编译并安装 custom_rasterizer
        whl_file_path = install_custom_rasterizer()

        if whl_file_path:
            print(f"安装完成，路径是: {whl_file_path}")
        else:
            print("安装 custom_rasterizer 失败，未生成 .whl 文件")
    
    except Exception as e:
        print(f"程序执行过程中发生错误: {str(e)}")