"""StarCoderData code data streaming.

Dataset: https://huggingface.co/datasets/bigcode/starcoderdata
Storage: ~50 GB for sampled subset
License: Permissive (varies by language)
"""
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 StarCoderConfig:
    languages: tuple = ("python", "javascript", "java", "cpp", "rust")
    ctx: int = 2048
    batch_size: int = 8
    shuffle_buffer: int = 10000
    split: str = "train"
    sample_rate: float = 0.05


class StarCoderStream:
    def __init__(self, cfg: StarCoderConfig):
        self.cfg = cfg
        self._ds = None
        self._bos = SPECIAL_VOCAB['BOS']
        self._eos = SPECIAL_VOCAB['EOS']
        self._pad = SPECIAL_VOCAB['PAD']

    def _lazy_init(self):
        if self._ds is not None:
            return
        from datasets import load_dataset, concatenate_datasets
        streams = []
        for lang in self.cfg.languages:
            ds = load_dataset("bigcode/starcoderdata", lang,
                              split=self.cfg.split, streaming=True)
            streams.append(ds)
        merged = concatenate_datasets(streams)
        self._ds = merged.shuffle(buffer_size=self.cfg.shuffle_buffer, seed=42)

    def _tokenize(self, text: str) -> torch.LongTensor:
        raw = text.encode("utf-8")
        tokens = [self._bos] + list(raw) + [self._eos]
        return torch.tensor(tokens[:self.cfg.ctx + 1], dtype=torch.long)

    def batches(self):
        self._lazy_init()
        buf = []
        for example in self._ds:
            tokens = self._tokenize(example["content"])
            if tokens.numel() < 5:
                continue
            buf.append(tokens)
            if len(buf) >= self.cfg.batch_size:
                batch = buf[:self.cfg.batch_size]
                buf = buf[self.cfg.batch_size:]
                max_len = max(t.numel() for t in batch)
                padded = torch.stack([
                    torch.cat([t, t.new_full((max_len - t.numel(),), self._pad)])
                    for t in batch
                ])
                yield padded[:, :-1].contiguous(), padded[:, 1:].contiguous()

    def num_samples(self) -> int:
        return 250_000_000_000 // self.cfg.ctx