import argparse
import os
import zipfile
from pathlib import Path


def extract_zip_to_images(dname):
    """
    Разархивирует {dname}.zip в папку images.
    """
    # Определяем пути
    current_dir = Path(__file__).parent
    zip_path = current_dir / f"{dname}.zip"
    extract_dir = current_dir

    # Проверяем существование архива
    if not zip_path.exists():
        print(f"Ошибка: архив не найден: {zip_path}")
        return False

    # Создаем папку для извлечения, если она не существует
    extract_dir.mkdir(exist_ok=True)

    try:
        # Открываем архив
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            # Получаем список файлов в архиве
            file_list = zip_ref.namelist()
            print(f"Найдено {len(file_list)} файлов в архиве")

            # Извлекаем все файлы
            zip_ref.extractall(extract_dir)

            # Получаем список извлеченных файлов
            extracted_files = list(extract_dir.rglob("*"))
            image_files = [f for f in extracted_files if f.is_file()]

            print(f"Успешно извлечено {len(image_files)} файлов в {extract_dir}")
            return True

    except zipfile.BadZipFile:
        print(f"Ошибка: архив поврежден или не является ZIP-файлом: {zip_path}")
        return False
    except Exception as e:
        print(f"Ошибка при разархивировании: {e}")
        return False


def main():
    """
    Основная функция для запуска разархивирования.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--dname', type=str, default='ya_go_maps', help='Dataset name')
    args = parser.parse_args()
    
    print(f"Начинаю разархивирование {args.dname}.zip...")

    if extract_zip_to_images(args.dname):
        print("Разархивирование успешно завершено!")
    else:
        print("Разархивирование не удалось.")


if __name__ == "__main__":
    main()