import os
import random
import shutil

from PIL import Image


def copy_convert_to_webp(
    source_dir,
    dest_dir,
    samples_per_folder=5,
    size_threshold=1024 * 1024,
    quality_high=90,
    quality_low=70,
):
    """
    Копирует структуру папок и конвертирует изображения в WebP
    :param source_dir: Исходная директория с изображениями
    :param dest_dir: Целевая директория для копирования
    :param samples_per_folder: Количество изображений для выбора из каждой папки
    :param size_threshold: Пороговый размер файла (в байтах) для сильного сжатия
    :param quality_high: Качество для маленьких изображений (0-100)
    :param quality_low: Качество для больших изображений (0-100)
    """
    image_extensions = (".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp")

    def is_image_file(filename):
        return filename.lower().endswith(image_extensions)

    def convert_to_webp(src_path, dest_path, quality):
        try:
            with Image.open(src_path) as img:
                # Сохраняем прозрачность для RGBA
                if img.mode in ("RGBA", "LA"):
                    img = img.convert("RGBA")
                    img.save(dest_path, "WEBP", quality=quality, lossless=False, method=6)
                else:
                    # Конвертируем в RGB для JPEG-подобных изображений
                    img = img.convert("RGB")
                    img.save(dest_path, "WEBP", quality=quality, method=6)
        except Exception as e:
            print(f"Error converting {src_path}: {e}")
            return False
        return True

    # Создаем корневую целевую папку
    os.makedirs(dest_dir, exist_ok=True)

    for root, dirs, files in os.walk(source_dir):
        rel_path = os.path.relpath(root, source_dir)
        target_dir = os.path.join(dest_dir, rel_path)
        os.makedirs(target_dir, exist_ok=True)

        images = [f for f in files if is_image_file(f)]
        if not images:
            continue

        selected = random.sample(images, min(samples_per_folder, len(images)))

        for file in selected:
            src_path = os.path.join(root, file)
            base_name = os.path.splitext(file)[0]
            dest_path = os.path.join(target_dir, f"{base_name}.webp")

            # Определяем качество на основе размера
            file_size = os.path.getsize(src_path)
            quality = quality_low if file_size > size_threshold else quality_high

            # Конвертируем в WebP
            if not convert_to_webp(src_path, dest_path, quality):
                # Если ошибка конвертации - копируем оригинал
                shutil.copy2(src_path, dest_path)
                print(f"Copied original file instead: {file}")


if __name__ == "__main__":
    copy_convert_to_webp(
        source_dir="assets",
        dest_dir="webp_output",
        samples_per_folder=5,
        size_threshold=1 * 1024 * 20,  # 1 MB
        quality_high=70,
        quality_low=70,
    )