#!/usr/bin/env python3
import argparse
import gzip
import json
import os
import re
import time
from collections import Counter, defaultdict
from pathlib import Path

from datasets import load_dataset
from transformers import AutoTokenizer


BUDGETS = {
    "english_web": 250_000_000,
    "english_edu": 200_000_000,
    "chinese_clean": 250_000_000,
    "code": 150_000_000,
    "math": 100_000_000,
    "science": 30_000_000,
    "qa_as_text": 20_000_000,
}


STREAM_SOURCES = {
    "english_web": [
        {"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-26", "split": "train", "max_rows": 0},
        {"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-21", "split": "train", "max_rows": 0},
    ],
    "english_edu": [
        {"kind": "hf", "name": "HuggingFaceFW/fineweb-edu", "config": None, "split": "train", "max_rows": 0},
    ],
    "chinese_clean": [
        {"kind": "hf", "name": "BAAI/CCI3-HQ", "config": None, "split": "train", "max_rows": 0},
        {"kind": "hf", "name": "Skywork/SkyPile-150B", "config": None, "split": "train", "max_rows": 0},
    ],
    "code": [
        {"kind": "hf", "name": "bigcode/starcoderdata", "config": None, "split": "train", "max_rows": 0},
        {"kind": "hf", "name": "codeparrot/github-code", "config": None, "split": "train", "max_rows": 0},
    ],
    "math": [
        {"kind": "hf", "name": "open-web-math/open-web-math", "config": None, "split": "train", "max_rows": 0},
        {"kind": "hf", "name": "GAIR/MathPile", "config": None, "split": "train", "max_rows": 0},
    ],
}


LOCAL_GLOBS = {
    "science": [
        "data/offline_text_only_reasoning_sources_20260611/science_reasoning__*.jsonl",
        "data/offline_text_only_reasoning_sources_20260611/logic__*.jsonl",
    ],
    "qa_as_text": [
        "data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl",
        "data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl",
    ],
}

LOCAL_PARQUET_GLOBS = {
    "english_web": [
        "data/raw_parquets/fineweb_2025/*.parquet",
    ],
    "english_edu": [
        "data/raw_parquets/fineweb_edu/*.parquet",
    ],
    "code": [
        "data/raw_parquets/starcoder/python/*.parquet",
        "data/raw_parquets/starcoder/javascript/*.parquet",
        "data/raw_parquets/starcoder/typescript/*.parquet",
        "data/raw_parquets/starcoder/java/*.parquet",
        "data/raw_parquets/starcoder/cpp/*.parquet",
        "data/raw_parquets/starcoder/go/*.parquet",
        "data/raw_parquets/starcoder/rust/*.parquet",
        "data/raw_parquets/starcoder/shell/*.parquet",
    ],
}


def clean_text(text):
    if text is None:
        return ""
    text = str(text).replace("\x00", " ")
    text = re.sub(r"[ \t\r\f\v]+", " ", text)
    text = re.sub(r"\n{4,}", "\n\n\n", text)
    return text.strip()


def first_user_assistant(messages):
    user = None
    assistant = None
    for msg in messages or []:
        role = msg.get("role")
        content = clean_text(msg.get("content"))
        if not content:
            continue
        if role == "user" and user is None:
            user = content
        elif role == "assistant" and user is not None:
            assistant = content
            break
    if user and assistant:
        return user, assistant
    return None, None


def qa_text(user, assistant):
    if not user or not assistant:
        return ""
    user_label = "" if user.lstrip().lower().startswith(("question:", "problem:", "context:", "support:", "fact1:")) else "Question:\n"
    assistant_label = "" if assistant.lstrip().lower().startswith(("answer:", "solution:", "final answer:", "explanation:")) else "Answer:\n"
    return f"{user_label}{user}\n\n{assistant_label}{assistant}"


def extract_text(row, category, source_name):
    if category == "science":
        user, assistant = first_user_assistant(row.get("messages"))
        if user and assistant:
            return qa_text(user, assistant)

    if category == "qa_as_text" or row.get("messages"):
        user, assistant = first_user_assistant(row.get("messages"))
        if user and assistant:
            return qa_text(user, assistant)

    if category == "code":
        content = clean_text(row.get("content") or row.get("text") or row.get("code"))
        if not content:
            return ""
        path = clean_text(row.get("path") or row.get("max_stars_repo_path") or row.get("repo_name") or "")
        lang = clean_text(row.get("lang") or row.get("language") or "")
        if path or lang:
            attrs = []
            if path:
                attrs.append(f'path="{path[:300]}"')
            if lang:
                attrs.append(f'language="{lang[:80]}"')
            return f"<file {' '.join(attrs)}>\n{content}\n</file>"
        return content

    for key in ("text", "content", "markdown", "raw_content"):
        text = clean_text(row.get(key))
        if text:
            return text
    return ""


def token_count(tok, text):
    return len(tok.encode(text, add_special_tokens=False))


def open_gz_writer(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    return gzip.open(path, "at", encoding="utf-8")


def write_doc(writer, doc_id, category, source, text, ntok, metadata=None):
    writer.write(
        json.dumps(
            {
                "id": doc_id,
                "category": category,
                "source": source,
                "text": text,
                "token_count": ntok,
                "metadata": metadata or {},
            },
            ensure_ascii=False,
        )
        + "\n"
    )


def should_keep(row, category, text, ntok, args):
    if ntok < args.min_tokens:
        return False, "too_short"
    if ntok > args.max_doc_tokens:
        return False, "too_long"
    if category == "english_web":
        lang = row.get("language")
        score = row.get("language_score")
        if lang and str(lang).lower() != "en":
            return False, "non_en"
        if score is not None:
            try:
                if float(score) < args.min_fineweb_language_score:
                    return False, "low_language_score"
            except Exception:
                pass
    return True, ""


def iter_local_jsonl(paths):
    for path in paths:
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    yield path, json.loads(line)


def flush_pack(writer, category, source, parts, ntok, stats, args):
    if not parts:
        return [], 0
    if ntok < args.min_tokens:
        stats["rejected"][category]["pack_too_short"] += 1
        return [], 0
    idx = stats["docs_by_category"][category]
    text = "\n\n---\n\n".join(parts)
    write_doc(
        writer,
        f"{category}_packed_{idx:09d}",
        category,
        source,
        text,
        ntok,
        {"local_path": source, "packed_items": len(parts)},
    )
    stats["docs_by_category"][category] += 1
    stats["tokens_by_category"][category] += ntok
    stats["tokens_by_source"][source] += ntok
    if stats["docs_by_category"][category] % args.log_every == 0:
        print_progress(stats, category)
    return [], 0


def collect_local(base, tok, category, out_dir, args, stats):
    paths = []
    for pattern in LOCAL_GLOBS.get(category, []):
        paths.extend(str(p) for p in sorted((base).glob(pattern)))
    if not paths:
        stats["sources"].append({"category": category, "kind": "local", "error": "no local files"})
        return

    out_path = out_dir / "documents" / f"{category}.jsonl.gz"
    with open_gz_writer(out_path) as writer:
        pack_parts = []
        pack_tokens = 0
        pack_source = ""
        for path, row in iter_local_jsonl(paths):
            if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                break
            text = extract_text(row, category, path)
            if not text:
                stats["rejected"][category]["empty"] += 1
                continue
            ntok = token_count(tok, text)
            if category == "science" and ntok < args.min_tokens:
                if ntok > args.max_doc_tokens:
                    stats["rejected"][category]["too_long"] += 1
                    continue
                if pack_source and pack_source != path:
                    pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
                pack_source = path
                if pack_tokens + ntok > args.science_pack_tokens and pack_parts:
                    pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
                pack_parts.append(text)
                pack_tokens += ntok
                continue
            keep, reason = should_keep(row, category, text, ntok, args)
            if not keep:
                stats["rejected"][category][reason] += 1
                continue
            if category == "science" and pack_parts:
                pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
            idx = stats["docs_by_category"][category]
            write_doc(writer, f"{category}_local_{idx:09d}", category, path, text, ntok, {"local_path": path})
            stats["docs_by_category"][category] += 1
            stats["tokens_by_category"][category] += ntok
            stats["tokens_by_source"][path] += ntok
            if stats["docs_by_category"][category] % args.log_every == 0:
                print_progress(stats, category)
        if category == "science" and pack_parts and stats["tokens_by_category"][category] < BUDGETS[category] * args.scale:
            flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)


def collect_local_parquet(base, tok, category, out_dir, args, stats):
    try:
        import pyarrow.parquet as pq
    except Exception as exc:
        stats["sources"].append({"category": category, "kind": "local_parquet", "error": repr(exc)})
        return

    paths = []
    for pattern in LOCAL_PARQUET_GLOBS.get(category, []):
        paths.extend(sorted(base.glob(pattern)))
    if not paths:
        stats["sources"].append({"category": category, "kind": "local_parquet", "error": "no parquet files"})
        return

    out_path = out_dir / "documents" / f"{category}.jsonl.gz"
    with open_gz_writer(out_path) as writer:
        for path in paths:
            if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                break
            source_label = str(path.relative_to(base))
            source_rec = {"category": category, "kind": "local_parquet", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
            try:
                pf = pq.ParquetFile(path)
                for batch in pf.iter_batches(batch_size=args.parquet_batch_size):
                    rows = batch.to_pylist()
                    for row in rows:
                        source_rec["rows_seen"] += 1
                        if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                            break
                        text = extract_text(row, category, source_label)
                        if not text:
                            stats["rejected"][category]["empty"] += 1
                            continue
                        ntok = token_count(tok, text)
                        keep, reason = should_keep(row, category, text, ntok, args)
                        if not keep:
                            stats["rejected"][category][reason] += 1
                            continue
                        idx = stats["docs_by_category"][category]
                        write_doc(
                            writer,
                            f"{category}_{safe_source(source_label)}_{idx:09d}",
                            category,
                            source_label,
                            text,
                            ntok,
                            {k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
                        )
                        stats["docs_by_category"][category] += 1
                        stats["tokens_by_category"][category] += ntok
                        stats["tokens_by_source"][source_label] += ntok
                        source_rec["docs_written"] += 1
                        source_rec["tokens"] += ntok
                        if stats["docs_by_category"][category] % args.log_every == 0:
                            print_progress(stats, category)
                    if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                        break
            except Exception as exc:
                source_rec["error"] = repr(exc)
                print(json.dumps({"event": "local_parquet_error", **source_rec}, ensure_ascii=False), flush=True)
            stats["sources"].append(source_rec)


def collect_stream(tok, category, out_dir, args, stats):
    out_path = out_dir / "documents" / f"{category}.jsonl.gz"
    with open_gz_writer(out_path) as writer:
        for spec in STREAM_SOURCES.get(category, []):
            if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                break
            source_label = spec["name"] if not spec.get("config") else f"{spec['name']}:{spec['config']}"
            source_rec = {"category": category, "kind": "hf_stream", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
            try:
                ds = load_dataset(
                    spec["name"],
                    spec.get("config"),
                    split=spec.get("split", "train"),
                    streaming=True,
                    trust_remote_code=True,
                )
                for row in ds:
                    source_rec["rows_seen"] += 1
                    if spec.get("max_rows") and source_rec["rows_seen"] > spec["max_rows"]:
                        break
                    if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
                        break
                    text = extract_text(row, category, source_label)
                    if not text:
                        stats["rejected"][category]["empty"] += 1
                        continue
                    ntok = token_count(tok, text)
                    keep, reason = should_keep(row, category, text, ntok, args)
                    if not keep:
                        stats["rejected"][category][reason] += 1
                        continue
                    idx = stats["docs_by_category"][category]
                    write_doc(
                        writer,
                        f"{category}_{safe_source(source_label)}_{idx:09d}",
                        category,
                        source_label,
                        text,
                        ntok,
                        {k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
                    )
                    stats["docs_by_category"][category] += 1
                    stats["tokens_by_category"][category] += ntok
                    stats["tokens_by_source"][source_label] += ntok
                    source_rec["docs_written"] += 1
                    source_rec["tokens"] += ntok
                    if stats["docs_by_category"][category] % args.log_every == 0:
                        print_progress(stats, category)
            except Exception as exc:
                source_rec["error"] = repr(exc)
                print(json.dumps({"event": "source_error", **source_rec}, ensure_ascii=False), flush=True)
            stats["sources"].append(source_rec)


def safe_source(text):
    return re.sub(r"[^A-Za-z0-9._-]+", "_", text)[:120]


def print_progress(stats, category):
    print(
        json.dumps(
            {
                "event": "progress",
                "category": category,
                "docs": stats["docs_by_category"][category],
                "tokens": stats["tokens_by_category"][category],
                "target": int(BUDGETS[category] * stats["scale"]),
                "elapsed_sec": time.time() - stats["start_time"],
            },
            ensure_ascii=False,
        ),
        flush=True,
    )


def dump_stats(out_dir, stats):
    serializable = {
        "scale": stats["scale"],
        "budgets": {k: int(v * stats["scale"]) for k, v in BUDGETS.items()},
        "tokens_by_category": dict(stats["tokens_by_category"]),
        "docs_by_category": dict(stats["docs_by_category"]),
        "tokens_by_source": dict(stats["tokens_by_source"].most_common()),
        "rejected": {k: dict(v) for k, v in stats["rejected"].items()},
        "sources": stats["sources"],
        "elapsed_sec": time.time() - stats["start_time"],
    }
    (out_dir / "stats.json").write_text(json.dumps(serializable, ensure_ascii=False, indent=2), encoding="utf-8")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
    parser.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
    parser.add_argument("--out-dir", default="data/cpt_docmix_1b_8192_20260613")
    parser.add_argument("--scale", type=float, default=1.0, help="1.0 = 1B token budget; 0.001 = 1M token smoke")
    parser.add_argument("--categories", default="english_web,english_edu,chinese_clean,code,math,science,qa_as_text")
    parser.add_argument("--min-tokens", type=int, default=128)
    parser.add_argument("--max-doc-tokens", type=int, default=32768)
    parser.add_argument("--min-fineweb-language-score", type=float, default=0.65)
    parser.add_argument("--log-every", type=int, default=1000)
    parser.add_argument("--parquet-batch-size", type=int, default=1000)
    parser.add_argument("--science-pack-tokens", type=int, default=2048)
    args = parser.parse_args()

    base = Path(args.base_dir)
    out_dir = base / args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "manifest.json").write_text(
        json.dumps(
            {
                "tokenizer": str(base / args.tokenizer),
                "budgets": {k: int(v * args.scale) for k, v in BUDGETS.items()},
                "seq_len_for_later_packing": 8192,
                "stream_sources": STREAM_SOURCES,
                "local_globs": LOCAL_GLOBS,
                "local_parquet_globs": LOCAL_PARQUET_GLOBS,
                "min_tokens": args.min_tokens,
                "max_doc_tokens": args.max_doc_tokens,
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

    tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
    stats = {
        "scale": args.scale,
        "start_time": time.time(),
        "tokens_by_category": Counter(),
        "docs_by_category": Counter(),
        "tokens_by_source": Counter(),
        "rejected": defaultdict(Counter),
        "sources": [],
    }

    for category in [x.strip() for x in args.categories.split(",") if x.strip()]:
        if category in LOCAL_PARQUET_GLOBS:
            collect_local_parquet(base, tok, category, out_dir, args, stats)
        if category in STREAM_SOURCES:
            collect_stream(tok, category, out_dir, args, stats)
        if category in LOCAL_GLOBS:
            collect_local(base, tok, category, out_dir, args, stats)
        dump_stats(out_dir, stats)

    dump_stats(out_dir, stats)
    print(out_dir / "stats.json")


if __name__ == "__main__":
    main()