import os
import json

def merge_json_by_config(input_root, config_name, output_file):
    """
    Merge all JSON files from each task's given config subfolder into a single list.
    
    Parameters:
    - input_root (str): path to the main folder containing task subfolders.
    - config_name (str): name of the config subfolder to merge from (e.g., 'demo_clean').
    - output_file (str): path to save the merged output JSON.
    """
    merged_data = []

    for task_name in sorted(os.listdir(input_root)):
        task_path = os.path.join(input_root, task_name)
        config_path = os.path.join(task_path, config_name)

        if os.path.isdir(config_path):
            for filename in sorted(os.listdir(config_path)):
                if filename.endswith(".json"):
                    file_path = os.path.join(config_path, filename)
                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            content = json.load(f)
                            if isinstance(content, list):
                                merged_data.extend(content)
                            else:
                                print(f"Warning: {file_path} does not contain a list. Skipping.")
                    except json.JSONDecodeError as e:
                        print(f"Warning: Failed to parse {file_path}: {e}")
        else:
            print(f"Warning: Config folder '{config_name}' not found in {task_name}. Skipping.")

    with open(output_file, "w", encoding="utf-8") as f_out:
        json.dump(merged_data, f_out, indent=2, ensure_ascii=False)

    print(f"✅ Merge complete. Output saved to: {output_file}")


input_folder = "yourpath/RoboTwin/policy/LLaVA-VLA/training_data"
config_name = "demo_clean" #or demo_randomized
output_file = "yourpath/RoboTwin/policy/LLaVA-VLA/training_data/data.json"  # Desired output file
merge_json_by_config(input_folder, config_name, output_file)