import json
from pathlib import Path


def generate_folder_structure(root_path):
    result = {}
    # Поддерживаемые форматы изображений
    image_ext = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff"}

    # Основные категории (cat и dog)
    for category in ["cat", "dog"]:
        category_path = Path(root_path) / category
        if not category_path.is_dir():
            continue

        category_dict = {}
        # Обрабатываем подпапки внутри категории
        for subfolder in category_path.iterdir():
            if subfolder.is_dir():
                # Собираем изображения
                images = [
                    file.name for file in subfolder.iterdir() if file.is_file() and file.suffix.lower() in image_ext
                ]
                category_dict[subfolder.name] = sorted(images)

        result[category] = category_dict

    return result


def save_to_json(data, output_file):
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    # Настройки
    ASSETS_DIR = "server/static/assets"
    OUTPUT_JSON = "structure.json"

    # Генерация структуры
    structure = generate_folder_structure(ASSETS_DIR)

    # Сохранение в файл
    save_to_json(structure, OUTPUT_JSON)
    print(f"JSON структура сохранена в файл: {OUTPUT_JSON}")