import os
import json

# Base directory where the ROOT files are located
base_dir = "/afs/crc.nd.edu/user/a/abasnet/topcoffea_photon/topeft/analysis/topeft_run2"
output_file = 'file_sizes.txt'

def get_root_file_size(root_file_path):
    try:
        size = os.path.getsize(root_file_path)
        return size
    except FileNotFoundError:
        print(f"File not found: {root_file_path}")
        return 0

def get_total_size(cfg_file, output):
    total_size = 0
    with open(cfg_file, 'r') as cfg:
        for line in cfg:
            json_file_path = line.strip().lstrip('#').strip()  # Handle commented lines
            if json_file_path.endswith('.json'):
                try:
                    with open(json_file_path, 'r') as json_file:
                        data = json.load(json_file)
                        root_files = data.get('files', [])
                        for root_file in root_files:
                            root_file_path = os.path.join(base_dir, root_file)
                            root_file_path = '/cms/cephfs/data'+root_file_path
                            file_size = get_root_file_size(root_file_path)
                            total_size += file_size
                            output.write(f"{root_file_path}: {file_size / (1024**4):.6f} TB\n")
                except FileNotFoundError:
                    print(f"JSON file not found: {json_file_path}")
                except json.JSONDecodeError:
                    print(f"Error decoding JSON file: {json_file_path}")
    return total_size

def calculate_total_size_across_cfgs(cfg_of_interest_file, output_file):
    total_size = 0
    with open(cfg_of_interest_file, 'r') as file:
        cfg_files = [line.strip() for line in file.readlines()]

    with open(output_file, 'w') as output:
        for cfg_file_path in cfg_files:
            output.write(f"Processing {cfg_file_path}...\n")
            total_size += get_total_size(cfg_file_path, output)
            output.write("\n")

    return total_size

# Path to the cfg_of_interest.txt file
cfg_of_interest_file = 'cfg_of_interest.txt'

total_size = calculate_total_size_across_cfgs(cfg_of_interest_file, output_file)
print(f"Total size of all ROOT files: {total_size / (1024**4):.2f} TB")