from functools import partial
from io import BytesIO
import pathlib
from typing import Any

import datasets
from loguru import logger
from PIL import Image
import requests
from tqdm.auto import tqdm

from src import config


def _save_resized_image(example: dict[str, Any], size: tuple[int, int], path: pathlib.Path):
    # Download the image
    image_url = example["url"]
    image_path = path / image_url.rsplit("/", 1)[-1]
    if image_path.exists():
        return

    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content))
    # Resize the image
    image_resized = image.resize(size)
    image_resized.save(image_path)


def _get_images(dataset: datasets.Dataset, path: pathlib.Path):
    save_resized_image = partial(_save_resized_image, path=path, size=(256, 256))
    dataset.map(save_resized_image, num_proc=128)


def _check_corrupt_images(image_file: pathlib.Path):
    try:
        with Image.open(image_file) as img:
            img.verify()  # Verify the integrity of the image
    except (IOError, SyntaxError) as e:
        logger.error(f"Corrupt image: {image_file}")


if __name__ == "__main__":
    hyper_parameters = config.TrainerConfig()

    dataset = datasets.load_dataset(
        hyper_parameters._data_config.dataset,
        split="train",
    )

    config.IMAGE_DOWNLOAD_PATH.mkdir(parents=True, exist_ok=True)
    _get_images(dataset, config.IMAGE_DOWNLOAD_PATH)  # type: ignore

    for image in tqdm(config.IMAGE_DOWNLOAD_PATH.iterdir()):
        _check_corrupt_images(image)