import os
import shutil
from datasets import load_from_disk, concatenate_datasets

def combine_datasets(main_dataset_path, datasets_to_add):
    """
    Объединяет указанные датасеты с основным датасетом.
    
    Args:
        main_dataset_path (str): Путь к основному датасету, в который нужно добавить данные
        datasets_to_add (list): Список путей к датасетам, которые нужно добавить
    
    Returns:
        Dataset: Объединенный датасет
    """
    # Загружаем основной датасет
    try:
        main_dataset = load_from_disk(main_dataset_path)
        print(f"Загружен основной датасет: {main_dataset_path} ({len(main_dataset)} записей)")
    except Exception as e:
        print(f"Ошибка загрузки основного датасета: {e}")
        return None
    
    # Список всех датасетов для объединения
    all_datasets = [main_dataset]
    
    # Загружаем и добавляем все дополнительные датасеты
    for path in datasets_to_add:
        try:
            ds = load_from_disk(path)
            all_datasets.append(ds)
            print(f"Добавлен датасет: {path} ({len(ds)} записей)")
        except Exception as e:
            print(f"Ошибка загрузки датасета {path}: {e}")
    
    # Объединяем все датасеты
    print(f"Объединение {len(all_datasets)} датасетов...")
    combined = concatenate_datasets(all_datasets)
    
    # Создаем временную директорию на основе имени основного датасета
    temp_dir = f"{main_dataset_path}_temp"
    
    # Удаляем временную директорию, если она уже существует
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    
    try:
        # Сохраняем в временную директорию
        print(f"Сохранение во временную директорию {temp_dir}...")
        combined.save_to_disk(temp_dir)
        
        # Удаляем старую директорию и перемещаем новую на ее место
        print(f"Обновление основного датасета...")
        #if os.path.exists(main_dataset_path):
        #    shutil.rmtree(main_dataset_path)
        #shutil.copytree(temp_dir, main_dataset_path)
        
        # Удаляем временную директорию после успешного копирования
        #shutil.rmtree(temp_dir)
        
        print(f"✅ Объединенный датасет ({len(combined)} записей) успешно сохранен в: {main_dataset_path}")
    except Exception as e:
        print(f"Ошибка при сохранении датасета: {e}")
        print(f"Временные данные сохранены в: {temp_dir}")
    
    return combined

combine_datasets("/workspace/sdxs3d/datasets/batch1_640", ["/workspace/sdxs3d/datasets/batch0_640_temp"])