from torch.utils.data import Dataset
import numpy as np
import torch
import lmdb
import json
from pathlib import Path
from PIL import Image
import os
import datasets

class TextDataset(Dataset):

    def __init__(self, prompt_path, extended_prompt_path=None):
        with open(prompt_path, encoding='utf-8') as f:
            self.prompt_list = [line.rstrip() for line in f]
        if extended_prompt_path is not None:
            with open(extended_prompt_path, encoding='utf-8') as f:
                self.extended_prompt_list = [line.rstrip() for line in f]
            assert len(self.extended_prompt_list) == len(self.prompt_list)
        else:
            self.extended_prompt_list = None

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

    def __getitem__(self, idx):
        batch = {'prompts': self.prompt_list[idx], 'idx': idx}
        if self.extended_prompt_list is not None:
            batch['extended_prompts'] = self.extended_prompt_list[idx]
        return batch

class TwoTextDataset(Dataset):

    def __init__(self, prompt_path: str, switch_prompt_path: str):
        with open(prompt_path, encoding='utf-8') as f:
            self.prompt_list = [line.rstrip() for line in f]
        with open(switch_prompt_path, encoding='utf-8') as f:
            self.switch_prompt_list = [line.rstrip() for line in f]
        assert len(self.switch_prompt_list) == len(self.prompt_list), 'The two prompt files must contain the same number of lines so that each first-segment prompt is paired with exactly one second-segment prompt.'

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

    def __getitem__(self, idx):
        return {'prompts': self.prompt_list[idx], 'switch_prompts': self.switch_prompt_list[idx], 'idx': idx}

class MultiTextDataset(Dataset):

    def __init__(self, prompt_path: str, field: str='prompts', cache_dir: str | None=None):
        self.ds = datasets.load_dataset('json', data_files=prompt_path, split='train', cache_dir=cache_dir, streaming=False)
        assert len(self.ds) > 0, 'JSONL is empty'
        assert field in self.ds.column_names, f"Missing field '{field}'"
        seg_len = len(self.ds[0][field])
        for i, ex in enumerate(self.ds):
            val = ex[field]
            assert isinstance(val, list), f"Line {i} field '{field}' is not a list"
            assert len(val) == seg_len, f'Line {i} list length mismatch'
        self.field = field

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

    def __getitem__(self, idx: int):
        return {'idx': idx, 'prompts_list': self.ds[idx][self.field]}

def cycle(dl):
    while True:
        for data in dl:
            yield data