import io
import os
import tarfile
from functools import lru_cache
from typing import Optional
from threading import Lock
from src.logger import get_logger

import pandas as pd
from PIL import Image


logger = get_logger(__name__)

_shard_locks = {}
_shard_locks_mutex = Lock()


@lru_cache(maxsize=64)
def extract_tarfile(shard_path: str):
    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 open_tarfile(shard_path: str):
    with _shard_locks_mutex:
        if shard_path not in _shard_locks:
            _shard_locks[shard_path] = Lock()
        lock = _shard_locks[shard_path]

    with lock:
        return extract_tarfile(shard_path)


@lru_cache(maxsize=8)
def load_index(index_path: str) -> dict[str, str]:
    df = pd.read_json(index_path, lines=True)
    index = dict(zip(df["image"], df["shard"]))
    return index


class ShardedImagesRepository:
    """
    A repository to manage and retrieve images stored in shards.

    The file structure should be as follows:
    (base_dir)/
        (source)/
            00001.tar
            00002.tar
            ...
            index.jsonl

    where each .tar file contains images and index.jsonl is "image" to "shard" mapping.
    """

    def __init__(self, base_dir: str) -> None:
        self.base_dir = base_dir

    def _read_from_shard(self, source: str, shard: str, image: str) -> Image.Image:
        members = open_tarfile(os.path.join(self.base_dir, source, shard))
        images_file = members.get(image, None)
        assert images_file is not None, f"Image {image} not found in shard {shard}"

        image = Image.open(io.BytesIO(images_file))
        return image

    def get_image(self, source: str, image: str) -> Optional[Image.Image]:
        source_index = load_index(os.path.join(self.base_dir, source, "index.jsonl"))
        if not source_index:
            return None

        shard = source_index.get(image)
        if not shard:
            return None

        return self._read_from_shard(source, shard, image)
