import os
import subprocess
import shutil
import argparse

# ==========================================
# 設定區
# ==========================================
BASE_DIR = "."

# 自動偵測作業系統，決定 Blender 的執行檔路徑
if os.name == 'nt':  # Windows 系統 (本機)
    BLENDER_EXE = os.path.join(BASE_DIR, "blender-3.6.0-windows-x64", "blender.exe")
else:                # Linux 系統 (Hugging Face Docker)
    BLENDER_EXE = os.path.join(BASE_DIR, "blender", "blender")

def run_cmd(cmd, cwd=BASE_DIR):
    """印出指令並執行"""
    print(f"\n[執行指令]: {' '.join(cmd)}")
    env = os.environ.copy()
    env['PYTHONPATH'] = BASE_DIR + os.pathsep + env.get('PYTHONPATH', '')
    subprocess.run(cmd, cwd=cwd, env=env, check=True)

def main():
    parser = argparse.ArgumentParser(description="Process a single video for CHAMP pipeline")
    parser.add_argument("--video", type=str, default="input_video/sample.mp4", help="Path to input video")
    parser.add_argument("--ref", type=str, default="reference_imgs/images/sample.jpg", help="Path to reference image")
    args = parser.parse_args()

    INPUT_VIDEO = args.video
    REF_IMAGE = args.ref

    # ---------------------------------------------------------
    # 0. 解析檔案名稱與路徑
    # ---------------------------------------------------------
    video_basename = os.path.basename(INPUT_VIDEO)
    video_name = os.path.splitext(video_basename)[0]
    
    ref_basename = os.path.basename(REF_IMAGE)
    ref_name = os.path.splitext(ref_basename)[0]
    
    driving_dir = os.path.join("driving_videos", video_name)
    images_dir = os.path.join(driving_dir, "images")
    smpl_results_dir = os.path.join(driving_dir, "smpl_results")
    
    # 參考圖的父資料夾 (例如: reference_imgs)
    ref_dir = os.path.dirname(os.path.dirname(REF_IMAGE)) 
    ref_npy_path = os.path.join(ref_dir, "smpl_results", f"{ref_name}.npy")
    
    transfer_dir = os.path.join("transferd_result", video_name)
    
    print("==================================================")
    print(f"正在處理單一影片: {video_name}")
    print(f"使用的參考圖片: {ref_name}")
    print("==================================================")

    # ---------------------------------------------------------
    # 1. 抽取影格 (ffmpeg)
    # ---------------------------------------------------------
    print("\n[步驟 1] 影片抽幀 (轉換為 24fps)...")
    images_abs_path = os.path.join(BASE_DIR, images_dir)
    os.makedirs(images_abs_path, exist_ok=True)
    out_pattern = os.path.join(images_abs_path, "%04d.png")
    
    if len(os.listdir(images_abs_path)) == 0:
        run_cmd([
            "ffmpeg", "-y", "-i", INPUT_VIDEO,
            "-vf", "fps=24,scale=640:-1",
            out_pattern
        ])
    else:
        print("影格圖片已存在，跳過抽幀。")

    # ---------------------------------------------------------
    # 2. 生成 SMPL (fit SMPL)
    # ---------------------------------------------------------
    print("\n[步驟 2] 為動作影片與參考圖生成 SMPL 骨架...")
    smpls_group_path = os.path.join(smpl_results_dir, "smpls_group.npz")
    if not os.path.exists(os.path.join(BASE_DIR, smpls_group_path)):
        run_cmd([
            "python", "-m", "scripts.data_processors.smpl.generate_smpls",
            "--reference_imgs_folder", ref_dir,
            "--driving_video_path", driving_dir
        ])
    else:
        print("smpls_group.npz 已存在，跳過生成。")

    # ---------------------------------------------------------
    # 3. 平滑化 SMPL (Smooth SMPL)
    # ---------------------------------------------------------
    print("\n[步驟 3] 使用 Blender 平滑化 SMPL 參數...")
    run_cmd([
        BLENDER_EXE, "--background",
        "--python", "scripts/data_processors/smpl/smooth_smpls.py",
        "--",
        "--smpls_group_path", smpls_group_path,
        "--smoothed_result_path", smpls_group_path
    ])

    # ---------------------------------------------------------
    # 4. 轉換 SMPL (Transfer SMPL)
    # ---------------------------------------------------------
    print("\n[步驟 4] 融合參考圖身型與影片動作 (Transfer SMPL)...")
    run_cmd([
        "python", "-m", "scripts.data_processors.smpl.smpl_transfer",
        "--reference_path", ref_npy_path,
        "--driving_path", driving_dir,
        "--output_folder", transfer_dir,
        "--figure_transfer",
        "--view_transfer"
    ])

    # ---------------------------------------------------------
    # 5. 渲染條件圖 (Render SMPL via Blender)
    # ---------------------------------------------------------
    print("\n[步驟 5] 渲染 3D 條件影像 (Depth, Normal, Mask 等)...")
    transfer_smpl_results = os.path.join(transfer_dir, "smpl_results")
    # 斷點續傳檢查：如果已經有足夠的 depth 圖片就不重跑 Blender
    depth_dir = os.path.join(BASE_DIR, transfer_dir, "depth")
    skip_blender = False
    if os.path.exists(depth_dir) and len(os.listdir(depth_dir)) > 0:
        print("渲染圖片似乎已存在，Blender 將會自動跳過已完成的部分。")
        
    run_cmd([
        BLENDER_EXE, "scripts/data_processors/smpl/blend/smpl_rendering.blend",
        "--background",
        "--python", "scripts/data_processors/smpl/render_condition_maps.py",
        "--",
        "--driving_path", transfer_smpl_results,
        "--reference_path", REF_IMAGE,
        "--device", "0"
    ])

    # ---------------------------------------------------------
    # 6. 生成 DWPose 骨架圖
    # ---------------------------------------------------------
    print("\n[步驟 6] 生成 DWPose 2D 骨架特徵圖...")
    transfer_normal = os.path.join(transfer_dir, "normal")
    transfer_dwpose = os.path.join(transfer_dir, "dwpose")
    run_cmd([
        "python", "-m", "scripts.data_processors.dwpose.generate_dwpose",
        "--input", transfer_normal,
        "--output", transfer_dwpose
    ])

    print("\n 完成 motions 轉換")

if __name__ == "__main__":
    main()