import os
import argparse
import subprocess
import sys

# python 2_Batch_Gen_All_PLY.py --z-scale 0.7 # --img-path "./Data" # --brightness 1.0  --saturation 1.0  --gamma 1.0

# === 配置区域 ===
# 默认图片路径，与 universal 脚本保持一致
DEFAULT_IMG_PATH = "./Data"
# 点云生成脚本的文件名 (确保该文件在当前目录下，且已经是修改支持颜色调整的版本)
PLY_SCRIPT_NAME = "Tool_Gen_3d_points_Cloud.py"
# ================

def run_ply_generation(ori_img_path, depth_folder_path, ply_out_path, args):
    """
    调用 Tool_Gen_3d_points_Cloud.py 生成点云，并传递颜色/形状调整参数
    """
    if not os.path.exists(PLY_SCRIPT_NAME):
        print(f"Error: 找不到脚本 {PLY_SCRIPT_NAME}，请确保它在当前目录下。")
        return False

    cmd = [
        sys.executable, PLY_SCRIPT_NAME,
        '--img_path_ori', ori_img_path,
        '--img_path_depth', depth_folder_path,
        '--outdir', ply_out_path,
        # 关键参数：universal脚本生成的图片没有 "_depth" 后缀，所以这里设为空字符串
        '--appendix', "", 
        '--label-type', 'Depth', # 默认就是Depth
        
        # === 新增透传参数 ===
        '--z-scale', str(args.z_scale),
        '--brightness', str(args.brightness),
        '--saturation', str(args.saturation),
        '--gamma', str(args.gamma)
    ]

    print(f"\n>>> Processing: {os.path.basename(depth_folder_path)} -> {os.path.basename(ply_out_path)}")
    print(f"    (Params: Z-Scale={args.z_scale}, Bright={args.brightness}, Sat={args.saturation}, Gamma={args.gamma})")
    
    try:
        # 调用子进程
        subprocess.check_call(cmd)
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error generating PLY for {depth_folder_path}: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Batch Generate PLY from Universal Results')
    parser.add_argument('--img-path', type=str, default=DEFAULT_IMG_PATH, help='Path to original input images folder')
    
    # === 新增调整参数 (与 Tool_Gen 脚本对应) ===
    parser.add_argument("--z-scale", type=float, default=1.0, help="Z轴拉伸倍数，数值越大，模型越立体/修长")
    parser.add_argument("--brightness", type=float, default=1.0, help="亮度调整倍数 (例如 1.2 为变亮20%)")
    parser.add_argument("--saturation", type=float, default=1.0, help="饱和度调整倍数 (例如 1.2 为颜色更鲜艳)")
    parser.add_argument("--gamma", type=float, default=1.0, help="Gamma校正 (数值<1.0 提亮暗部细节)")
    # ==========================================
    
    args = parser.parse_args()

    # 1. 确定路径
    input_path_abs = os.path.abspath(args.img_path)
    
    # 推断 universal 脚本生成的总结果目录名
    # 规则是: 文件夹名 + "_Universal_Results"
    dir_basename = os.path.basename(os.path.normpath(input_path_abs))
    root_result_dir = os.path.abspath(f"{dir_basename}_Universal_Results")

    if not os.path.exists(root_result_dir):
        print(f"Error: 找不到结果目录: {root_result_dir}")
        print("请先运行 '1_Batch_Gen_All_Depth(跑之前删除结果文件夹).py' 生成深度图。")
        return

    print(f"Original Images: {input_path_abs}")
    print(f"Scanning Results in: {root_result_dir}")

    # 2. 扫描该目录下所有的文件夹
    subfolders = [f for f in os.listdir(root_result_dir) if os.path.isdir(os.path.join(root_result_dir, f))]
    
    task_count = 0

    for folder_name in sorted(subfolders):
        # 只处理灰度深度图文件夹 (_gray_result)
        if folder_name.endswith("_gray_result"):
            
            # 构建源路径 (深度图)
            depth_in_path = os.path.join(root_result_dir, folder_name)
            
            # 构建目标路径 (PLY结果)
            # 将 _gray_result 替换为 _ply_result
            ply_folder_name = folder_name.replace("_gray_result", "_ply_result")
            ply_out_path = os.path.join(root_result_dir, ply_folder_name)

            # 3. 执行转换 (传入 args)
            success = run_ply_generation(input_path_abs, depth_in_path, ply_out_path, args)
            if success:
                task_count += 1

    if task_count == 0:
        print("\n未找到任何 '_gray_result' 文件夹。请检查上一以步是否生成了灰度图。")
    else:
        print(f"\n=== 全部完成，共处理了 {task_count} 组模型数据 ===")
        print(f"结果保存在: {root_result_dir} 下的 *_ply_result 文件夹中")

if __name__ == "__main__":
    main()