import os
import csv

# Root directory to process
ROOT_DIR = "../data/a1.0.0_00"
OUTPUT_CSV = "../result/s1_output.csv"

rows = []
for folder_name in os.listdir(ROOT_DIR):
    folder_path = os.path.join(ROOT_DIR, folder_name)
    if os.path.isdir(folder_path):
        # Collect all .png files except 'assembly.png'
        png_files = [
            f for f in os.listdir(folder_path)
            if f.endswith(".png") and f != "assembly.png"
        ]
        # Remove '.png' extension from file names
        clean_names = [os.path.splitext(f)[0] for f in png_files]
        # Join all names with ';'
        joined_names = ";".join(sorted(clean_names))
        count = len(clean_names)
        # Append to CSV rows
        rows.append([folder_name, joined_names, count])

with open(OUTPUT_CSV, mode="w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["Assembly", "Part_names", "Part_count"])  # CSV header
    writer.writerows(rows)

print(f"CSV saved to {OUTPUT_CSV}")
