#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from multiprocessing import Pool
from pathlib import Path
from types import SimpleNamespace
from typing import Iterable

import pyarrow.parquet as pq

_TOKENIZER = None
_APPEND_EOD = True
_MAX_SEQ_LEN = 65536


def add_megatron_to_path(megatron_dir: str) -> None:
    path = str(Path(megatron_dir).resolve())
    if path not in sys.path:
        sys.path.insert(0, path)


def make_tokenizer_args(args: argparse.Namespace) -> SimpleNamespace:
    return SimpleNamespace(
        rank=0,
        make_vocab_size_divisible_by=128,
        tensor_model_parallel_size=1,
        pad_vocab_size=True,
        padded_vocab_size=None,
        vocab_size=args.vocab_size,
        vocab_file=args.vocab_file,
        merge_file=args.merge_file,
        vocab_extra_ids=0,
        tokenizer_type=args.tokenizer_type,
        tokenizer_model=args.tokenizer_model,
        metadata_path=args.tokenizer_metadata,
        special_tokens=args.tokenizer_special_tokens,
        tokenizer_sentencepiece_legacy=args.tokenizer_sentencepiece_legacy,
        tokenizer_hf_no_use_fast=args.tokenizer_hf_no_use_fast,
        tokenizer_hf_no_include_special_tokens=args.tokenizer_hf_no_include_special_tokens,
        trust_remote_code=args.trust_remote_code,
        tiktoken_pattern=args.tiktoken_pattern,
        tiktoken_num_special_tokens=args.tiktoken_num_special_tokens,
        null_tokenizer_eod_id=args.null_tokenizer_eod_id,
        null_tokenizer_pad_id=args.null_tokenizer_pad_id,
        tokenizer_prompt_format=None,
        chat_template=None,
        image_tag_type=None,
        force_system_message=False,
        sft_tokenizer_prompt_format=None,
    )


def build_tokenizer(args: argparse.Namespace):
    add_megatron_to_path(args.megatron_dir)
    from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer

    return build_tokenizer(make_tokenizer_args(args))


def init_worker(args: argparse.Namespace) -> None:
    global _TOKENIZER, _APPEND_EOD, _MAX_SEQ_LEN
    _APPEND_EOD = args.append_eod
    _MAX_SEQ_LEN = args.max_seq_len
    _TOKENIZER = build_tokenizer(args)
    if _APPEND_EOD and _TOKENIZER.eod is None:
        raise ValueError("Tokenizer has no EOD/EOS token, but --append-eod is enabled.")


def encode_text(text: object):
    if not isinstance(text, str):
        return None
    text = text.strip()
    if not text:
        return None
    token_ids = _TOKENIZER.tokenize(text)
    if not token_ids:
        return None

    if _MAX_SEQ_LEN and _MAX_SEQ_LEN > 0:
        content_chunk_size = _MAX_SEQ_LEN - 1 if _APPEND_EOD else _MAX_SEQ_LEN
        chunks = []
        for start in range(0, len(token_ids), content_chunk_size):
            chunk = token_ids[start : start + content_chunk_size]
            if _APPEND_EOD:
                chunk = [*chunk, _TOKENIZER.eod]
            chunks.append((chunk, [len(chunk)]))
        return chunks

    if _APPEND_EOD:
        token_ids.append(_TOKENIZER.eod)
    return [(token_ids, [len(token_ids)])]


def output_paths(output_prefix: Path, text_key: str) -> tuple[Path, Path]:
    return (
        Path(str(output_prefix) + f"_{text_key}_document.bin"),
        Path(str(output_prefix) + f"_{text_key}_document.idx"),
    )


def resolve_inputs(inputs: Iterable[str], max_files: int | None) -> list[Path]:
    paths: list[Path] = []
    for item in inputs:
        path = Path(item)
        if path.is_dir():
            paths.extend(path.glob("*.parquet"))
            paths.extend(path.glob("*.zstd.parquet"))
            continue
        matches = sorted(Path().glob(item)) if not path.is_absolute() else sorted(path.parent.glob(path.name))
        if matches:
            paths.extend(matches)
            continue
        if path.exists():
            paths.append(path)
            continue
        raise FileNotFoundError(item)
    unique = sorted({p.resolve() for p in paths}, key=str)
    if max_files is not None:
        unique = unique[:max_files]
    return unique


def safe_prefix_component(value: str) -> str:
    value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
    value = value.strip("._-")
    return value or "source"


def convert_one_parquet(task: tuple[str, argparse.Namespace]) -> dict[str, object]:
    parquet_path = Path(task[0])
    args = task[1]
    add_megatron_to_path(args.megatron_dir)
    from megatron.core.datasets import indexed_dataset

    stem = parquet_path.name.replace(".zstd.parquet", "").replace(".parquet", "")
    source_name = safe_prefix_component(parquet_path.parent.name)
    output_prefix = Path(args.output_dir) / f"{args.output_prefix_prefix}_{source_name}_{stem}"
    bin_file, idx_file = output_paths(output_prefix, args.text_key)
    if not args.overwrite and bin_file.exists() and idx_file.exists():
        return {
            "status": "skip",
            "source": str(parquet_path),
            "output_prefix": str(output_prefix),
            "bin_file": str(bin_file),
            "idx_file": str(idx_file),
            "reason": "existing bin/idx",
        }

    bin_file.unlink(missing_ok=True)
    idx_file.unlink(missing_ok=True)
    output_prefix.parent.mkdir(parents=True, exist_ok=True)

    tokenizer = build_tokenizer(args)
    dtype = indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size)
    builder = indexed_dataset.IndexedDatasetBuilder(str(bin_file), dtype=dtype)

    start = time.time()
    rows = source_docs = output_docs = tokens = split_source_docs = split_extra_docs = 0
    stop = False

    def consume(encoded) -> None:
        nonlocal source_docs, output_docs, tokens, split_source_docs, split_extra_docs, stop
        if encoded is None or stop:
            return
        source_docs += 1
        if len(encoded) > 1:
            split_source_docs += 1
            split_extra_docs += len(encoded) - 1
        for token_ids, lengths in encoded:
            if args.max_seq_len and len(token_ids) > args.max_seq_len:
                raise ValueError(
                    f"internal split error: chunk has {len(token_ids)} tokens, "
                    f"max_seq_len={args.max_seq_len}"
                )
            builder.add_document(token_ids, lengths)
            output_docs += 1
            tokens += len(token_ids)
        if args.log_interval and source_docs % args.log_interval == 0:
            elapsed = max(time.time() - start, 1e-6)
            print(
                f"[{parquet_path.name}] source_docs={source_docs} output_docs={output_docs} "
                f"tokens={tokens} source_docs/s={source_docs / elapsed:.2f}",
                flush=True,
            )
        if args.max_docs and source_docs >= args.max_docs:
            stop = True

    pf = pq.ParquetFile(parquet_path)
    if args.text_key not in pf.schema_arrow.names:
        raise ValueError(f"{parquet_path} has no column {args.text_key!r}")

    if args.workers_per_file == 1:
        init_worker(args)
        for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
            texts = batch.column(0).to_pylist()
            rows += len(texts)
            for text in texts:
                consume(encode_text(text))
                if stop:
                    break
            if stop:
                break
    else:
        with Pool(processes=args.workers_per_file, initializer=init_worker, initargs=(args,)) as pool:
            for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
                texts = batch.column(0).to_pylist()
                rows += len(texts)
                for encoded in pool.imap(encode_text, texts, chunksize=args.chunksize):
                    consume(encoded)
                    if stop:
                        break
                if stop:
                    break

    builder.finalize(str(idx_file))
    elapsed = max(time.time() - start, 1e-6)
    return {
        "status": "ok",
        "source": str(parquet_path),
        "output_prefix": str(output_prefix),
        "bin_file": str(bin_file),
        "idx_file": str(idx_file),
        "dtype": dtype.__name__,
        "rows": rows,
        "source_docs": source_docs,
        "docs": output_docs,
        "output_docs": output_docs,
        "split_source_docs": split_source_docs,
        "split_extra_docs": split_extra_docs,
        "max_seq_len": args.max_seq_len,
        "tokens": tokens,
        "elapsed_sec": elapsed,
        "source_docs_per_sec": source_docs / elapsed,
        "docs_per_sec": output_docs / elapsed,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Convert pretrain parquet directly to Megatron .bin/.idx.")
    parser.add_argument("--input", action="append", required=True, help="Parquet dir/file/glob. Repeatable.")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--manifest", default=None)
    parser.add_argument("--megatron-dir", default="/opt/Megatron-Bridge/3rdparty/Megatron-LM")
    parser.add_argument("--tokenizer-type", default="HuggingFaceTokenizer")
    parser.add_argument("--tokenizer-model", required=True)
    parser.add_argument("--tokenizer-metadata", default=None)
    parser.add_argument("--tokenizer-special-tokens", nargs="*", default=None)
    parser.add_argument("--tokenizer-sentencepiece-legacy", action="store_true")
    parser.add_argument("--tokenizer-hf-no-use-fast", action="store_true")
    parser.add_argument("--tokenizer-hf-no-include-special-tokens", action="store_true")
    parser.add_argument("--trust-remote-code", action="store_true")
    parser.add_argument("--vocab-file", default=None)
    parser.add_argument("--merge-file", default=None)
    parser.add_argument("--vocab-size", type=int, default=None)
    parser.add_argument("--tiktoken-pattern", default=None)
    parser.add_argument("--tiktoken-num-special-tokens", type=int, default=1000)
    parser.add_argument("--null-tokenizer-eod-id", type=int, default=None)
    parser.add_argument("--null-tokenizer-pad-id", type=int, default=-1)
    parser.add_argument("--text-key", default="text")
    parser.add_argument("--output-prefix-prefix", default="laoyao_2b_moe_8192")
    parser.add_argument("--parallel-files", type=int, default=1)
    parser.add_argument("--workers-per-file", type=int, default=max((os.cpu_count() or 8) // 4, 1))
    parser.add_argument("--batch-size", type=int, default=8192)
    parser.add_argument("--chunksize", type=int, default=128)
    parser.add_argument("--max-files", type=int, default=None)
    parser.add_argument("--max-docs", type=int, default=0, help="Per-file doc cap; 0 means no cap.")
    parser.add_argument("--max-seq-len", type=int, default=65536, help="Split documents into chunks no longer than this; 0 disables splitting.")
    parser.add_argument("--log-interval", type=int, default=10000)
    parser.add_argument("--append-eod", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.parallel_files < 1 or args.workers_per_file < 1:
        raise ValueError("parallel and worker counts must be positive")
    if args.max_seq_len < 0:
        raise ValueError("--max-seq-len must be >= 0")
    if args.append_eod and args.max_seq_len == 1:
        raise ValueError("--max-seq-len must be > 1 when --append-eod is enabled")
    files = resolve_inputs(args.input, args.max_files)
    if not files:
        raise FileNotFoundError("no parquet files resolved")

    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    manifest_path = Path(args.manifest) if args.manifest else Path(args.output_dir) / "manifest.json"
    print(
        json.dumps(
            {
                "files": len(files),
                "parallel_files": args.parallel_files,
                "workers_per_file": args.workers_per_file,
                "total_tokenizer_workers": args.parallel_files * args.workers_per_file,
                "output_dir": args.output_dir,
                "max_seq_len": args.max_seq_len,
            },
            ensure_ascii=False,
        ),
        flush=True,
    )

    tasks = [(str(path), args) for path in files]
    results: list[dict[str, object]] = []
    if args.parallel_files == 1:
        for task in tasks:
            result = convert_one_parquet(task)
            results.append(result)
            print(json.dumps(result, ensure_ascii=False), flush=True)
    else:
        with ProcessPoolExecutor(max_workers=args.parallel_files) as executor:
            futures = [executor.submit(convert_one_parquet, task) for task in tasks]
            for future in as_completed(futures):
                result = future.result()
                results.append(result)
                print(json.dumps(result, ensure_ascii=False), flush=True)

    manifest = {
        "args": vars(args),
        "files": [str(path) for path in files],
        "results": sorted(results, key=lambda item: str(item.get("source", ""))),
        "total_source_docs": sum(int(r.get("source_docs", 0)) for r in results),
        "total_docs": sum(int(r.get("docs", 0)) for r in results),
        "total_output_docs": sum(int(r.get("output_docs", 0)) for r in results),
        "total_split_source_docs": sum(int(r.get("split_source_docs", 0)) for r in results),
        "total_split_extra_docs": sum(int(r.get("split_extra_docs", 0)) for r in results),
        "total_tokens": sum(int(r.get("tokens", 0)) for r in results),
        "ok_prefixes": [
            str(Path(str(r["bin_file"])).with_suffix(""))
            for r in results
            if r.get("status") in {"ok", "skip"}
        ],
    }
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"manifest={manifest_path}", flush=True)


if __name__ == "__main__":
    main()
