from functools import lru_cache
import io
import os
import tarfile
from typing import Union, Optional
from PIL import Image


@lru_cache(maxsize=64)
def open_tarfile(shard_path):
    members = {}
    with tarfile.open(shard_path, "r") as tar:
        for member in tar.getmembers():
            if member is not None and member.isfile():
                members[member.name] = tar.extractfile(member).read()

    return members


def read_image(root: str, name: str, shard: Optional[str] = None) -> Image.Image:
    if shard:
        shard_path = os.path.join(root, shard)
        members = open_tarfile(shard_path)
        images_file = members.get(name, None)
        assert images_file is not None, f"Image {name} not found in shard {shard_path}"
        image = Image.open(io.BytesIO(images_file))
        return image
    else:
        image_path = os.path.join(root, name)
        image = Image.open(image_path)
        return image


def fetch_image(url: str) -> Image.Image:
    import requests

    response = requests.get(url)
    response.raise_for_status()
    image = Image.open(io.BytesIO(response.content))
    return image


def verify(image: Union[Image.Image, bytes]) -> bool:
    try:
        if isinstance(image, Image.Image):
            image.verify()
            return True
        if isinstance(image, bytes):
            image = Image.open(io.BytesIO(image))
            image.verify()
            return True

        return False
    except Exception:
        return False
