import hashlib
import os
import pickle

import torch
from torch.utils.data import Dataset
from torchvision.transforms import CenterCrop, Normalize, Resize
from torchvision.transforms.functional import to_tensor
from PIL import Image

EXTs = ['.png', '.jpg', '.jpeg', ".JPEG"]


def is_image_file(filename):
    return any(filename.endswith(ext) for ext in EXTs)


def _index_root(root: str, cache_path: str = None, require_text: bool = True):
    """Walk ``root`` once and return a list of (image_path, text_path) pairs.

    The result is cached on disk so subsequent epochs / DDP workers do not
    repeat the walk.  Caching is keyed by ``root`` + mtime of ``root`` so the
    cache invalidates when the top-level directory changes.
    """
    if cache_path is None:
        key = hashlib.md5(root.encode("utf-8")).hexdigest()[:16]
        cache_path = os.path.join("/tmp", f"imagetext_index_{key}.pkl")

    try:
        st = os.stat(root)
        mtime = st.st_mtime
    except FileNotFoundError:
        raise FileNotFoundError(f"dataset root does not exist: {root}")

    if os.path.exists(cache_path):
        try:
            with open(cache_path, "rb") as f:
                cached = pickle.load(f)
            if cached.get("root") == root and cached.get("mtime") == mtime:
                return cached["pairs"]
        except Exception:
            pass

    pairs = []
    for dir_, _subdirs, files in os.walk(root):
        # split images from txts using a set so we don't stat per file
        files_set = set(files)
        for file in files:
            if not is_image_file(file):
                continue
            base, _ = os.path.splitext(file)
            txt_name = base + ".txt"
            if require_text and txt_name not in files_set:
                continue
            image_path = os.path.join(dir_, file)
            text_path = os.path.join(dir_, txt_name) if txt_name in files_set else None
            pairs.append((image_path, text_path))

    try:
        with open(cache_path, "wb") as f:
            pickle.dump({"root": root, "mtime": mtime, "pairs": pairs}, f,
                        protocol=pickle.HIGHEST_PROTOCOL)
    except Exception:
        pass

    return pairs


class ImageText(Dataset):
    """Image-text dataset.

    Behavior change vs. the original implementation: the ``__init__`` no longer
    opens and reads every ``.txt`` file -- it only indexes file *paths*.  The
    actual text is read on demand in ``__getitem__``.  An index of (image,
    text) path pairs is also cached on disk under ``/tmp`` so that subsequent
    runs / DDP worker processes start instantly.
    """

    def __init__(self, root: str, resolution: int, cache_path: str = None,
                 require_text: bool = True):
        super().__init__()
        self.root = root
        self.pairs = _index_root(root, cache_path=cache_path, require_text=require_text)

        self.resize = Resize(resolution)
        self.center_crop = CenterCrop(resolution)
        self.normalize = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

    @property
    def image_paths(self):  # backward-compat for any external user
        return [p[0] for p in self.pairs]

    def __getitem__(self, idx: int):
        image_path, text_path = self.pairs[idx]
        if text_path is not None:
            try:
                with open(text_path, "r") as f:
                    text = f.read()
            except Exception:
                text = ""
        else:
            text = ""

        pil_image = Image.open(image_path).convert('RGB')
        pil_image = self.resize(pil_image)
        pil_image = self.center_crop(pil_image)
        raw_image = to_tensor(pil_image)
        normalized_image = self.normalize(raw_image)
        metadata = {
            "image_path": image_path,
            "prompt": text,
            "raw_image": raw_image,
        }
        return normalized_image, text, metadata

    def __len__(self):
        return len(self.pairs)