"""CC12M image-text data preparation.

Dataset: https://huggingface.co/datasets/opendiffusionai/cc12m-4mp-realistic
Storage: ~250 GB for images
License: MIT (dataset), images vary
"""
import torch, io
from dataclasses import dataclass
from PIL import Image
from torchvision import transforms
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
from arbitor.config import SPECIAL_VOCAB


@dataclass
class CC12MConfig:
    batch_size: int = 4
    shuffle_buffer: int = 5000
    split: str = "train"
    image_size: int = 224
    max_samples: int = 0


class CC12MStream:
    def __init__(self, cfg: CC12MConfig):
        self.cfg = cfg
        self._ds = None
        self._bos = SPECIAL_VOCAB['BOS']
        self._eos = SPECIAL_VOCAB['EOS']
        self._pad = SPECIAL_VOCAB['PAD']
        self._image_tok = SPECIAL_VOCAB['IMAGE']
        self._transform = transforms.Compose([
            transforms.Resize((cfg.image_size, cfg.image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def _lazy_init(self):
        if self._ds is not None:
            return
        from datasets import load_dataset
        ds = load_dataset("opendiffusionai/cc12m-4mp-realistic",
                          split=self.cfg.split, streaming=True)
        if self.cfg.max_samples:
            ds = ds.take(self.cfg.max_samples)
        self._ds = ds.shuffle(buffer_size=self.cfg.shuffle_buffer, seed=42)

    def _decode_image(self, img_bytes):
        try:
            pil = Image.open(io.BytesIO(img_bytes)).convert("RGB")
            return self._transform(pil)
        except Exception:
            return None

    def batches(self):
        self._lazy_init()
        buf = []
        for example in self._ds:
            img_tensor = self._decode_image(example["image"]["bytes"])
            if img_tensor is None:
                continue
            caption = example.get("text", "")
            raw = caption.encode("utf-8")
            text_tokens = [self._bos] + list(raw) + [self._eos]
            text_t = torch.tensor(text_tokens, dtype=torch.long)
            buf.append((img_tensor, text_t))
            if len(buf) >= self.cfg.batch_size:
                batch = buf[:self.cfg.batch_size]
                buf = buf[self.cfg.batch_size:]
                imgs = torch.stack([b[0] for b in batch])
                max_len = max(b[1].numel() for b in batch)
                texts = torch.stack([
                    torch.cat([b[1], b[1].new_full((max_len - b[1].numel(),), self._pad)])
                    for b in batch
                ])
                yield imgs, texts

    def num_samples(self) -> int:
        return 12_000_000