import os
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import Dataset


def txt_dir_save_fn(image, metadata, root_path):
    stem = metadata["filename"]
    Image.fromarray(image).save(os.path.join(root_path, f"{stem}.png"))
    with open(os.path.join(root_path, f"{stem}.txt"), "w") as f:
        f.write(metadata["prompt"])


class TxtDirEvalDataset(Dataset):
    def __init__(
        self,
        txt_dir: str,
        latent_shape,
        limit: int = 50,
        seed_offset: int = 0,
    ):
        self.latent_shape = latent_shape
        self.seed_offset = seed_offset
        self.txt_paths = sorted(Path(txt_dir).glob("*.txt"))
        if limit is not None and limit > 0:
            self.txt_paths = self.txt_paths[:limit]
        if not self.txt_paths:
            raise FileNotFoundError(f"No .txt files found in {txt_dir}")

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

    def __getitem__(self, idx):
        path = self.txt_paths[idx]
        prompt = path.read_text().strip()
        seed = self.seed_offset + idx
        generator = torch.Generator().manual_seed(seed)
        latent = torch.randn(self.latent_shape, generator=generator, dtype=torch.float32)
        metadata = dict(
            prompt=prompt,
            filename=path.stem,
            seed=seed,
            save_fn=txt_dir_save_fn,
        )
        return latent, prompt, metadata