import os
import yaml
import subprocess
import time
import shutil

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"  [x] [{name}] Failed. Elapsed: {elapsed:.1f}s")
            return False
        print(f"  [*] [{name}] Completed. Elapsed: {elapsed:.1f}s")
        return True
    except Exception as e:
        print(f"  [x] [{name}] Exception: {e}")
        return False

def main():
    config_path = "/root/autodl-tmp/SplatAtlas/configs/3dgut_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/3dgrut:{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})")
            
            ckpt_path = os.path.join(model_out, f"ours_{iters}/ckpt_{iters}.pt")
            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(ckpt_path) or os.path.exists(os.path.join(model_out, "ckpt_last.pt")):
                print(f"  [*] [Checkpoint] Weights found.")
                if not os.path.exists(render_flag):
                    print(f"  [!] [Cleanup] Previous render incomplete, re-rendering...")
                    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:
                    print(f"  [*] [Checkpoint] Render assets complete.")
            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} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}"
                run_step("Metrics Calculation", eval_cmd, curr_env)
            else:
                print(f"  [*] [Checkpoint] Metrics JSON already exists.")

            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)
            
            print(f"[*] {scene} Done!\n")

if __name__ == "__main__":
    main()