from collections import Counter
from pathlib import Path

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

# One real source (wiki) and three fake sources; 0 = real, 1 = fake
# The same identity basename appears in every source - splitting must happen
# at the identity level to prevent leakage (see splits.py), e.g:
# data_dir/SOURCE/identity/BASENAME.jpg -> same BASENAME for all SOURCEs go into same split
SOURCES = {
    "wiki": 0,
    "inpainting": 1,
    "text2img": 1,
    "insight": 1,
}


# Extracts source name from path assuming data_dir/source/identity/image.jpg layout
def get_source_name(path: Path) -> str:
    return Path(path).parent.parent.name


# Walks data_dir/source/identity/*.jpg and collects (path, label) pairs
class DFFDataset(Dataset):
    def __init__(self, data_dir, sources=None, transform=None):
        self.transform = transform
        self.samples = []

        data_dir = Path(data_dir)
        if not data_dir.exists():
            raise FileNotFoundError(
                f"Dataset root not found: {data_dir}. Expected a directory containing "
                "wiki/, inpainting/, text2img/, and insight/."
            )
        if sources is None:
            sources = list(SOURCES.keys())

        for source in sources:
            label = SOURCES[source]
            source_dir = data_dir / source
            if not source_dir.exists():
                raise FileNotFoundError(
                    f"Missing source directory: {source_dir}. Check `data_dir` in the config."
                )
            for subdir in sorted(source_dir.iterdir()):
                if subdir.is_dir():
                    for img_path in sorted(subdir.glob("*.jpg")):
                        self.samples.append((img_path, label))

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

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        image = Image.open(path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, label

    # Useful for quickly verifying class balance before training
    def label_counts(self):
        return Counter(label for _, label in self.samples)


# Wraps a list of prediction records as a Dataset for re-scoring via DataLoader.
class PathDataset(Dataset):
    def __init__(self, records, image_size, preprocess=None):
        from src.preprocessing import get_transforms
        self.records = records
        self.transform = get_transforms(train=False, image_size=image_size)
        self.preprocess = preprocess

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

    def __getitem__(self, idx):
        record = self.records[idx]
        image = Image.open(record["path"]).convert("RGB")
        if self.preprocess is not None:
            image = self.preprocess(image)
        return self.transform(image), record["label"], idx
