import os
import yaml
import subprocess
import time
import shutil
import glob
import re

def run_step(name, cmd, env):
    print(f"  [Step] {name}...")
    start_time = time.time()
    try:
        result = subprocess.run(cmd, shell=True, env=env)
        elapsed = time.time() - start_time
        if result.returncode != 0:
            print(f"  [Failed] {name}. Elapsed: {elapsed:.1f}s")
            return False
        print(f"  [Completed] {name}. Elapsed: {elapsed:.1f}s")
        return True
    except Exception as e:
        print(f"  [Exception] {name}: {e}")
        return False

def main():
    config_path = "/root/autodl-tmp/SplatAtlas/configs/absgssgf_benchmark.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)

    method = cfg["method_name"]
    iters = cfg["global_settings"]["iterations"]

    curr_env = os.environ.copy()
    curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/absgs_official:{curr_env.get('PYTHONPATH', '')}"

    for ds in cfg["datasets"]:
        print(f"\n>>>> Dataset Domain: {ds['name']} <<<<")
        for scene in ds["scenes"]:
            source = os.path.join(ds["base_path"], scene)
            model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}"
            res = ds.get("resolution", 1)
            
            if not os.path.exists(source):
                continue

            print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})")
            
            ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply")
            metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json")
            render_flag = os.path.join(model_out, f"render_complete_{iters}.flag")
            renders_dir = os.path.join(model_out, f"renders_test_{iters}")
            depths_dir = os.path.join(model_out, f"depths_test_{iters}")
            
            if os.path.exists(ply_path):
                print(f"  [Checkpoint] Ply file found.")
                if not os.path.exists(render_flag):
                    for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]:
                        d_path = os.path.join(model_out, f"{d}_{iters}")
                        if os.path.exists(d_path): shutil.rmtree(d_path)
                    
                    render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}"
                    if not run_step("Offline Rendering", render_cmd, curr_env): continue
            else:
                ckpt_list = glob.glob(os.path.join(model_out, "chkpnt*.pth"))
                if ckpt_list:
                    latest_ckpt = max(ckpt_list, key=lambda x: int(re.findall(r'chkpnt(\d+)\.pth', os.path.basename(x))[0] if re.findall(r'chkpnt(\d+)\.pth', os.path.basename(x)) else 0))
                    print(f"  [Resume] 发现断点，从 {latest_ckpt} 续训...")
                    train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling --start_checkpoint {latest_ckpt}"
                else:
                    train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling"
                
                if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue

            if not os.path.exists(metrics_json):
                eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}"
                run_step("Metrics Calculation", eval_cmd, curr_env)

            physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}"
            run_step("Offline Physics Diagnosis", physics_cmd, curr_env)

            run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env)
            
if __name__ == "__main__":
    main()