# This script computes the sizes of the folders listed in args.json

import os
import json
import csv
from tqdm import tqdm

class Folder:
    def __init__(self, folder_path, desc_prefix=""):
        self.folder_path = folder_path
        self.size = 0
        self.files = 0
        self.desc_prefix = desc_prefix
    
    def desc(self) -> str:
        return f"{self.desc_prefix}{self.folder_path} - {self.size}B"
    
    def get_size(self) -> int:
        """Return the size in Bytes of the folder"""
        p_bar = tqdm(desc=self.desc(), unit=' file')
        for path, dirs, files in os.walk(self.folder_path):
            for f in files:
                fp = os.path.join(path, f)
                self.size += os.path.getsize(fp)
                self.files += 1
                p_bar.set_description(self.desc(), refresh=False)
                p_bar.update()

        p_bar.set_description(self.desc(), refresh=True)
        p_bar.close()
        return self.size


if __name__ == "__main__":
    # Initialize CSV
    csv_file = open('output.csv', 'w', newline='')
    spamwriter = csv.writer(csv_file, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    spamwriter.writerow(['folder_path', 'size_B', 'files'])

    folders = []

    with open('args.json', 'r') as file:
        data = json.load(file)

    nb_folders = len(data['folder_paths'])

    for i, folder_path in enumerate(data['folder_paths']):
        folders.append(Folder(folder_path, desc_prefix=f"[{i+1: >3}/{nb_folders}] "))
    
    # Calculate the size of all folders
    for folder in folders:
        folder.get_size()
        # Write info in CSV
        spamwriter.writerow([folder.folder_path, folder.size, folder.files])
    
    csv_file.close()