"""WebVid-10M video data preparation with pig-vae latent encoding.

Dataset: https://huggingface.co/datasets/TempoFunk/webvid-10M
Storage: ~5 GB (CLIP features) or ~TB (raw video)
License: Research only

Note: For 450GB constraint, we download CLIP features or pre-encoded
latents rather than raw video files.
"""
import torch
from dataclasses import dataclass
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
from arbitor.config import SPECIAL_VOCAB


@dataclass
class WebVidConfig:
    batch_size: int = 2
    shuffle_buffer: int = 2000
    split: str = "train"
    max_samples: int = 0
    n_frames: int = 4
    latent_h: int = 32
    latent_w: int = 32


class WebVidStream:
    def __init__(self, cfg: WebVidConfig):
        self.cfg = cfg
        self._ds = None
        self._vae = None
        self._bos = SPECIAL_VOCAB['BOS']
        self._eos = SPECIAL_VOCAB['EOS']
        self._pad = SPECIAL_VOCAB['PAD']
        self._video_tok = SPECIAL_VOCAB['VIDEO']

    def _lazy_init(self):
        if self._ds is not None:
            return
        from datasets import load_dataset
        ds = load_dataset("TempoFunk/webvid-10M",
                          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 _get_vae(self):
        if self._vae is None:
            from arbitor.encoders.pig_vae import load_vae
            self._vae = load_vae(device="cuda")
            self._vae.eval()
            for p in self._vae.parameters():
                p.requires_grad = False
        return self._vae

    def batches(self):
        self._lazy_init()
        buf = []
        for example in self._ds:
            caption = example.get("text") or example.get("caption", "")
            raw = caption.encode("utf-8")
            text_tokens = [self._bos] + list(raw) + [self._eos]
            text_t = torch.tensor(text_tokens, dtype=torch.long)

            video = example.get("video")
            if video is not None:
                frames = torch.tensor(video["array"], dtype=torch.float32)
                if frames.dim() == 4:
                    frames = frames.permute(3, 0, 1, 2)
                if frames.shape[1] > self.cfg.n_frames:
                    idx = torch.linspace(0, frames.shape[1] - 1, self.cfg.n_frames).long()
                    frames = frames[:, idx]
                frames = frames.unsqueeze(0).cuda()
                vae = self._get_vae()
                with torch.no_grad():
                    latents = vae.encode(frames).cpu()
            else:
                latents = torch.randn(1, 16, self.cfg.n_frames, self.cfg.latent_h, self.cfg.latent_w)

            buf.append((text_t, latents.squeeze(0)))
            if len(buf) >= self.cfg.batch_size:
                batch = buf[:self.cfg.batch_size]
                buf = buf[self.cfg.batch_size:]
                texts = [b[0] for b in batch]
                max_len = max(t.numel() for t in texts)
                text_batch = torch.stack([
                    torch.cat([t, t.new_full((max_len - t.numel(),), self._pad)])
                    for t in texts
                ])
                latent_batch = torch.stack([b[1] for b in batch])
                yield text_batch, latent_batch

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