from __future__ import annotations

import argparse
import json
import os
import random
from collections import Counter
from pathlib import Path

import torch
from tokenizers import Tokenizer
from torch.utils.data import Dataset, IterableDataset, get_worker_info


T5_TOKENIZER = "/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json"
TEXT_KEYS = ("text", "content", "document", "raw_text")


def load_tokenizer(path: str = T5_TOKENIZER):
    tok = Tokenizer.from_file(path)
    bos = tok.token_to_id("</s>")
    eos = tok.token_to_id("</s>")
    pad = tok.token_to_id("<pad>")
    return tok, int(bos), int(eos), int(pad)


def token_name(tok: Tokenizer, idx: int) -> str:
    return tok.id_to_token(int(idx)) or f"<id:{int(idx)}>"


class CachedTokenDataset(Dataset):
    def __init__(self, cache_path: str, tokenizer_path: str = T5_TOKENIZER) -> None:
        self.tokenizer, self.bos_id, self.eos_id, self.pad_id = load_tokenizer(tokenizer_path)
        cache = torch.load(cache_path, map_location="cpu")
        self.ids = cache["ids"].to(torch.int32)
        self.seen_count = int(cache.get("seen_count", 0))
        self.skipped_count = int(cache.get("skipped_count", 0))
        self.bos_id = int(cache.get("bos_id", self.bos_id))
        self.eos_id = int(cache.get("eos_id", self.eos_id))
        self.pad_id = int(cache.get("pad_id", self.pad_id))

    def __len__(self) -> int:
        return self.ids.size(0)

    def __getitem__(self, idx: int) -> torch.Tensor:
        return self.ids[idx].long()


def data_files(path: str) -> list[Path]:
    p = Path(path)
    if p.is_file():
        return [p]
    return sorted(x for x in p.rglob("*") if x.suffix.lower() in {".txt", ".jsonl", ".json", ".parquet"})


def iter_file_text(f: Path, text_column: str):
    if f.suffix == ".txt":
        for line in f.read_text(encoding="utf-8", errors="ignore").splitlines():
            if line.strip():
                yield line.strip()
    elif f.suffix in {".jsonl", ".json"}:
        for line in f.read_text(encoding="utf-8", errors="ignore").splitlines():
            if line.strip():
                obj = json.loads(line)
                yield obj.get(text_column) or next(obj[k] for k in TEXT_KEYS if k in obj)
    elif f.suffix == ".parquet":
        import pyarrow.parquet as pq

        pf = pq.ParquetFile(f)
        names = set(pf.schema.names)
        col = text_column if text_column in names else next(k for k in TEXT_KEYS if k in names)
        for batch in pf.iter_batches(columns=[col], batch_size=2048):
            for item in batch.column(0).to_pylist():
                if item:
                    yield item


def short(text: str, n: int = 80) -> str:
    return text.replace("\n", "\\n").replace("\t", "\\t")[:n]


def max_run_info(row: list[int], tok: Tokenizer) -> tuple[int, str]:
    best = cur = 1
    best_token = row[0] if row else 0
    for a, b in zip(row, row[1:]):
        if a == b:
            cur += 1
            if cur > best:
                best = cur
                best_token = b
        else:
            cur = 1
    return best, short(tok.decode([best_token], skip_special_tokens=False))


def repetitive_pattern_info(row: list[int], tok: Tokenizer, max_pattern_len: int = 16, max_repeats: int = 3) -> tuple[int, int, str]:
    n = len(row)
    for width in range(2, max_pattern_len + 1):
        i = 0
        while i + width * (max_repeats + 1) <= n:
            pattern = row[i : i + width]
            repeats = 1
            j = i + width
            while j + width <= n and row[j : j + width] == pattern:
                repeats += 1
                if repeats > max_repeats:
                    return repeats, width, short(tok.decode(pattern, skip_special_tokens=False))
                j += width
            i += 1
    return 0, 0, ""


def repeated_phrase_ngram_info(row: list[int], n: int, tok: Tokenizer, threshold: int = 16) -> tuple[int, str]:
    counts = Counter(zip(*(row[i:] for i in range(n))))
    for ngram, count in counts.most_common():
        if count < threshold:
            break
        text = tok.decode(list(ngram), skip_special_tokens=False).strip()
        if len(text) >= 3 and " " in text and any(ch.isalnum() for ch in text):
            return count, short(text)
    return 0, ""


def reject_reasons(row: list[int], unk_id: int, tok: Tokenizer) -> list[str]:
    reasons = []
    unique = len(set(row))
    if unique <= 48:
        reasons.append(f"token_unique<=48(unique={unique})")
    bigram_count, bigram_text = repeated_phrase_ngram_info(row, 2, tok)
    if bigram_count >= 16:
        reasons.append(f"bigram_repeat>=16(count={bigram_count},text={bigram_text})")
    trigram_count, trigram_text = repeated_phrase_ngram_info(row, 3, tok)
    if trigram_count >= 16:
        reasons.append(f"trigram_repeat>=16(count={trigram_count},text={trigram_text})")
    if unk_id in row:
        reasons.append(f"has_unk(token=<unk>,id={unk_id})")
    run_count, run_token = max_run_info(row, tok)
    if run_count > 10:
        reasons.append(f"single_token_run>10(count={run_count},token={run_token})")
    pattern_repeats, pattern_width, pattern_text = repetitive_pattern_info(row, tok)
    if pattern_repeats > 3:
        reasons.append(f"repetitive_pattern>3(repeats={pattern_repeats},width={pattern_width},text={pattern_text})")
    return reasons


def append_reject(path: str, line: str) -> None:
    import fcntl

    with open(path, "a", encoding="utf-8", buffering=1) as f:
        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
        f.write(line + "\n")
        fcntl.flock(f.fileno(), fcntl.LOCK_UN)


class OnlinePackedDataset(IterableDataset):
    def __init__(
        self,
        data_path: str,
        tokenizer_path: str = T5_TOKENIZER,
        text_column: str = "text",
        pack_len: int = 1023,
        append_eos: bool = True,
        shuffle_buffer: int = 8192,
        reject_txt: str = "cache/online_rejected.txt",
        seed: int = 1234,
        rank: int = 0,
        world: int = 1,
    ) -> None:
        self.tokenizer_path = tokenizer_path
        self.tokenizer, self.bos_id, self.eos_id, self.pad_id = load_tokenizer(tokenizer_path)
        self.unk_id = int(self.tokenizer.token_to_id("<unk>"))
        self.files = data_files(data_path)
        self.text_column = text_column
        self.pack_len = pack_len
        self.append_eos = append_eos
        self.shuffle_buffer = shuffle_buffer
        self.reject_txt = reject_txt
        self.seed = seed
        self.rank = rank
        self.world = world
        self.seen_count = 0
        self.skipped_count = 0
        Path(reject_txt).parent.mkdir(parents=True, exist_ok=True)
        if rank == 0:
            Path(reject_txt).write_text("", encoding="utf-8")

    def __iter__(self):
        info = get_worker_info()
        worker_id = info.id if info else 0
        num_workers = info.num_workers if info else 1
        rank_files = self.files[self.rank :: self.world]
        files = rank_files[worker_id::num_workers]
        if not files:
            return
        tok = Tokenizer.from_file(self.tokenizer_path)
        eos_id = int(tok.token_to_id("</s>"))
        unk_id = int(tok.token_to_id("<unk>"))
        rng = random.Random(self.seed + 1009 * self.rank + worker_id)
        stream: list[int] = []
        rows: list[torch.Tensor] = []
        seen = accepted = rejected = packed = 0

        while True:
            rng.shuffle(files)
            for file_name in files:
                for text in iter_file_text(file_name, self.text_column):
                    seen += 1
                    ids = [int(x) for x in tok.encode(text, add_special_tokens=False).ids]
                    reasons = reject_reasons(ids, unk_id, tok)
                    if reasons:
                        rejected += 1
                        clean_text = text.replace("\n", "\\n").replace("\t", "\\t")
                        append_reject(
                            self.reject_txt,
                            f"rank={self.rank}\tworker={worker_id}\tdoc={seen}\tntok={len(ids)}\treason={','.join(reasons)}\ttext={clean_text}",
                        )
                        continue

                    accepted += 1
                    stream.extend(ids)
                    stream.append(eos_id)
                    while len(stream) >= self.pack_len:
                        payload = stream[: self.pack_len]
                        del stream[: self.pack_len]
                        row = payload + ([eos_id] if self.append_eos else [])
                        rows.append(torch.tensor(row, dtype=torch.long))
                        packed += 1
                        if len(rows) >= self.shuffle_buffer:
                            i = rng.randrange(len(rows))
                            rows[i], rows[-1] = rows[-1], rows[i]
                            yield rows.pop()

                    if seen % 10000 == 0:
                        print(
                            f"[online data rank={self.rank} worker={worker_id}] docs={seen} accepted={accepted} rejected={rejected} rows={packed} buffered={len(stream)}",
                            flush=True,
                        )

            while rows:
                i = rng.randrange(len(rows))
                rows[i], rows[-1] = rows[-1], rows[i]
                yield rows.pop()


class ELFMultipartPackedDataset(IterableDataset):
    """Stream ELF's pre-tokenized OWT HFDS shards as direct fixed-length rows."""

    def __init__(
        self,
        data_path: str,
        tokenizer_path: str = T5_TOKENIZER,
        pack_len: int = 1023,
        append_eos: bool = True,
        shuffle_buffer: int = 8192,
        seed: int = 1234,
        rank: int = 0,
        world: int = 1,
    ) -> None:
        self.data_path = Path(data_path)
        self.tokenizer_path = tokenizer_path
        self.tokenizer, self.bos_id, self.eos_id, self.pad_id = load_tokenizer(tokenizer_path)
        self.max_len = int(pack_len) + int(bool(append_eos))
        self.shuffle_buffer = shuffle_buffer
        self.seed = seed
        self.rank = rank
        self.world = world
        self.seen_count = 0
        self.skipped_count = 0
        parts_root = self.data_path / "parts"
        if parts_root.exists():
            self.parts = sorted(x for x in parts_root.iterdir() if x.is_dir() and x.name.startswith("part-"))
            self.shard_rows = False
        else:
            self.parts = [self.data_path]
            self.shard_rows = True
        if not self.parts:
            raise FileNotFoundError(f"no ELF HFDS parts under {self.data_path}")

    def _to_fixed_row(self, ids) -> torch.Tensor:
        row = [int(x) for x in ids[: self.max_len]]
        if len(row) < self.max_len:
            row.extend([int(self.pad_id)] * (self.max_len - len(row)))
        return torch.tensor(row, dtype=torch.long)

    def __iter__(self):
        from datasets import Sequence, load_from_disk
        from datasets.features import features as hf_features

        # Some ELF/T5 caches were saved with HF metadata feature type "List".
        # Newer datasets builds may not register that alias, so keep old caches loadable.
        hf_features._FEATURE_TYPES.setdefault("List", Sequence)

        info = get_worker_info()
        worker_id = info.id if info else 0
        num_workers = info.num_workers if info else 1
        if self.shard_rows:
            # HF datasets saved as one root with data-*.arrow need row-level sharding.
            # Otherwise only rank0 would own the single "part" and all other ranks are empty.
            parts = self.parts
            shard_count = max(1, self.world * num_workers)
            shard_index = self.rank * num_workers + worker_id
        else:
            rank_parts = self.parts[self.rank :: self.world]
            parts = rank_parts[worker_id::num_workers]
            if not parts:
                return
            shard_count = 1
            shard_index = 0

        rng = random.Random(self.seed + 1009 * self.rank + worker_id)
        rows: list[torch.Tensor] = []
        seen = yielded = 0

        while True:
            rng.shuffle(parts)
            for part in parts:
                ds = load_from_disk(str(part))
                if self.shard_rows:
                    ds = ds.shard(num_shards=shard_count, index=shard_index, contiguous=True)
                for item in ds:
                    ids = item.get("input_ids")
                    if ids is None or len(ids) == 0:
                        continue
                    seen += 1
                    rows.append(self._to_fixed_row(ids))
                    if len(rows) >= self.shuffle_buffer:
                        i = rng.randrange(len(rows))
                        rows[i], rows[-1] = rows[-1], rows[i]
                        yielded += 1
                        yield rows.pop()
                    if seen % 100000 == 0:
                        print(
                            f"[elf_hfds direct rank={self.rank} worker={worker_id}] docs={seen} rows={yielded} buffered={len(rows)} part={part.name}",
                            flush=True,
                        )

            while rows:
                i = rng.randrange(len(rows))
                rows[i], rows[-1] = rows[-1], rows[i]
                yielded += 1
                yield rows.pop()


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--cache_path", required=True)
    p.add_argument("--tokenizer_path", default=T5_TOKENIZER)
    args = p.parse_args()

    ds = CachedTokenDataset(args.cache_path, args.tokenizer_path)
    tok = ds.tokenizer
    print(f"rows={len(ds)} length={ds.ids.size(1)} seen={ds.seen_count} dropped={ds.skipped_count}")
    print(f"bos={ds.bos_id}:{token_name(tok, ds.bos_id)} eos={ds.eos_id}:{token_name(tok, ds.eos_id)}")
    print("head", [token_name(tok, x) for x in ds.ids[0, :16].tolist()])
    print("tail", [token_name(tok, x) for x in ds.ids[0, -16:].tolist()])
    print(f"cache_size={Path(args.cache_path).stat().st_size / 1024**3:.2f} GiB")


if __name__ == "__main__":
    main()