"""
data/tokenize_extra.py — 대용량 korean_extra/ 데이터셋 병렬 토큰화

HuggingFace datasets disk 포맷(arrow), parquet, jsonl 등 세 가지 포맷을
자동 감지하여 SentencePiece 토크나이저로 토큰화하고, 결과를 uint16 memmap
(.bin) 파일로 저장한다.  881 GB 이상의 대용량 데이터도 스트리밍·청크 방식으로
처리한다.

출력 포맷은 data/dataset.py PackedDataset / TextDataset 과 완전히 호환되는
numpy uint16 플랫 배열이다.

사용 예시:
    # 단일 디렉토리
    python data/tokenize_extra.py \
        --input_dir data/korean_extra/fineweb2_edu_ko \
        --output    data/fineweb2_train.bin \
        --num_proc  8

    # korean_extra/ 전체 서브디렉토리 일괄 처리
    python data/tokenize_extra.py \
        --input_dir data/korean_extra \
        --auto_scan \
        --output_dir data \
        --num_proc 8

    # 공개 검증
    python -c "
    import numpy as np
    d = np.memmap('data/fineweb2_train.bin', dtype='uint16', mode='r')
    print(f'총 토큰: {len(d):,}')
    "
"""

from __future__ import annotations

import argparse
import json
import math
import multiprocessing as mp
import os
import struct
import sys
import time
from pathlib import Path
from typing import Generator, Iterable, Iterator

import numpy as np
from tqdm import tqdm

# ---------------------------------------------------------------------------
# SentencePiece 임포트 (선택적 — 없으면 오류 메시지 출력 후 종료)
# ---------------------------------------------------------------------------
try:
    import sentencepiece as spm
except ImportError:
    print(
        "ERROR: sentencepiece 패키지가 설치되어 있지 않습니다.\n"
        "       pip install sentencepiece  로 설치 후 재실행하세요.",
        file=sys.stderr,
    )
    sys.exit(1)

# ---------------------------------------------------------------------------
# datasets 임포트
# ---------------------------------------------------------------------------
try:
    import datasets as hf_datasets
except ImportError:
    print(
        "ERROR: datasets 패키지가 설치되어 있지 않습니다.\n"
        "       pip install datasets  로 설치 후 재실행하세요.",
        file=sys.stderr,
    )
    sys.exit(1)


# ===========================================================================
# 상수
# ===========================================================================

UINT16_MAX      = 65535           # uint16 오버플로 경계
MIN_TOKENS      = 100             # 최소 토큰 수 (미만이면 버림)
MAX_TOKENS      = 32_768          # 최대 토큰 수 (초과분은 버림)
HANGUL_RE_THRESHOLD = 0.10        # 한글 비율 최소 기준 (이 미만이고 한글 아닌 경우 버림)
CHUNK_TOKENS    = 500_000         # memmap 청크 단위 (tokens)
EOS_TOKEN_PLACEHOLDER = 1         # EOS id — SP 기본값, 실제 id는 모델에서 읽음

# ---------------------------------------------------------------------------
# 한글 비율 필터
# ---------------------------------------------------------------------------
# ord 범위: 가(AC00) ~ 힣(D7A3), ㄱ(3131) ~ ㅣ(3163)
_HANGUL_START = 0xAC00
_HANGUL_END   = 0xD7A3


def _has_enough_korean_or_english(text: str) -> bool:
    """
    한글 문자 비율이 HANGUL_RE_THRESHOLD 이상이거나,
    ASCII 알파벳 비율이 0.3 이상이면 True 반환.
    둘 다 아닌 경우 False (중국어, 일본어만 있는 등).
    """
    if not text:
        return False
    total = len(text)
    hangul_cnt = sum(1 for ch in text if _HANGUL_START <= ord(ch) <= _HANGUL_END)
    if hangul_cnt / total >= HANGUL_RE_THRESHOLD:
        return True
    ascii_alpha = sum(1 for ch in text if ch.isascii() and ch.isalpha())
    if ascii_alpha / total >= 0.30:
        return True
    return False


# ===========================================================================
# 토크나이저 래퍼 (프로세스 간 공유 불가 — 각 워커에서 reload)
# ===========================================================================

class SPTokenizer:
    """SentencePiece 모델을 wrapping한 간단한 토크나이저."""

    def __init__(self, model_path: str) -> None:
        self._model_path = model_path
        self._sp: spm.SentencePieceProcessor | None = None

    # 프로세스 fork 후 _sp가 None인 경우 lazy load
    def _ensure_loaded(self) -> None:
        if self._sp is None:
            sp = spm.SentencePieceProcessor()
            sp.Load(self._model_path)
            self._sp = sp

    @property
    def eos_id(self) -> int:
        self._ensure_loaded()
        return self._sp.eos_id()

    @property
    def vocab_size(self) -> int:
        self._ensure_loaded()
        return self._sp.GetPieceSize()

    def encode(self, text: str) -> list[int]:
        self._ensure_loaded()
        return self._sp.EncodeAsIds(text)


# ===========================================================================
# 포맷 감지 & 이터레이터
# ===========================================================================

def _detect_format(input_dir: Path) -> str:
    """
    디렉토리 내용을 보고 포맷을 자동 감지한다.

    반환값:
        "hf_arrow"  — HuggingFace datasets disk 포맷 (dataset_info.json 존재)
        "parquet"   — .parquet 파일이 있음
        "jsonl"     — .jsonl 또는 .json 파일이 있음
        "unknown"   — 알 수 없음
    """
    if not input_dir.is_dir():
        raise NotADirectoryError(f"입력 경로가 디렉토리가 아닙니다: {input_dir}")

    # HF arrow 포맷 판별 — dataset_info.json 또는 state.json이 있으면 HF 포맷
    if (input_dir / "dataset_info.json").exists():
        return "hf_arrow"
    if (input_dir / "state.json").exists():
        return "hf_arrow"
    # 서브 디렉토리 안에 dataset_info.json이 있는 경우 (split 포함)
    for child in input_dir.iterdir():
        if child.is_dir() and (child / "dataset_info.json").exists():
            return "hf_arrow"

    # parquet 파일 확인
    parquets = list(input_dir.rglob("*.parquet"))
    if parquets:
        return "parquet"

    # jsonl / json 파일 확인
    jsonls = list(input_dir.rglob("*.jsonl")) + list(input_dir.rglob("*.json"))
    if jsonls:
        return "jsonl"

    return "unknown"


def _iter_hf_arrow(
    input_dir: Path,
    text_col: str,
    num_proc: int,
) -> Iterator[str]:
    """HuggingFace datasets disk 포맷에서 텍스트를 스트리밍한다."""
    print(f"  [포맷] HuggingFace arrow (disk): {input_dir}")
    try:
        ds = hf_datasets.load_from_disk(str(input_dir))
    except Exception as exc:
        # DatasetDict일 수 있음 — 'train' split 시도
        try:
            ds_dict = hf_datasets.load_from_disk(str(input_dir))
            if isinstance(ds_dict, hf_datasets.DatasetDict):
                splits = list(ds_dict.keys())
                print(f"  DatasetDict 감지. splits={splits}, 'train' split 사용.")
                ds = ds_dict.get("train", ds_dict[splits[0]])
            else:
                raise exc
        except Exception:
            raise RuntimeError(
                f"HF arrow 포맷 로드 실패: {input_dir}\n원인: {exc}"
            ) from exc

    # 실제 텍스트 컬럼 이름 추정
    col = _resolve_text_col(list(ds.column_names), text_col)
    print(f"  텍스트 컬럼: '{col}', 총 행 수: {len(ds):,}")

    for row in ds:
        yield row[col]


def _iter_parquet(input_dir: Path, text_col: str) -> Iterator[str]:
    """parquet 파일에서 텍스트를 스트리밍한다."""
    try:
        import pyarrow.parquet as pq  # type: ignore
    except ImportError:
        # datasets로 fallback
        print("  [경고] pyarrow 미설치, datasets로 parquet 로드 시도...")
        files = sorted(input_dir.rglob("*.parquet"))
        print(f"  [포맷] parquet ({len(files)} 파일): {input_dir}")
        ds = hf_datasets.load_dataset(
            "parquet",
            data_files={"train": [str(f) for f in files]},
            split="train",
            streaming=True,
        )
        col = _resolve_text_col(list(ds.column_names), text_col)
        print(f"  텍스트 컬럼: '{col}'")
        for row in ds:
            yield row[col]
        return

    files = sorted(input_dir.rglob("*.parquet"))
    print(f"  [포맷] parquet ({len(files)} 파일): {input_dir}")
    for fpath in files:
        pf = pq.ParquetFile(str(fpath))
        cols = pf.schema_arrow.names
        col = _resolve_text_col(cols, text_col)
        for batch in pf.iter_batches(batch_size=1000, columns=[col]):
            for val in batch.column(col):
                yield val.as_py() or ""


def _iter_jsonl(input_dir: Path, text_col: str) -> Iterator[str]:
    """jsonl / json 파일에서 텍스트를 스트리밍한다."""
    files = sorted(input_dir.rglob("*.jsonl")) + sorted(input_dir.rglob("*.json"))
    # json 파일 중 jsonl이 아닌 것 제거 (파일 자체가 dict인 경우)
    print(f"  [포맷] jsonl ({len(files)} 파일): {input_dir}")
    for fpath in files:
        try:
            with open(fpath, "r", encoding="utf-8", errors="replace") as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        obj = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    if isinstance(obj, str):
                        yield obj
                    elif isinstance(obj, dict):
                        text = (
                            obj.get(text_col)
                            or obj.get("text")
                            or obj.get("content")
                            or obj.get("document")
                            or ""
                        )
                        yield str(text)
        except Exception as exc:
            print(f"  [경고] 파일 읽기 실패: {fpath} — {exc}", file=sys.stderr)


def _resolve_text_col(columns: list[str], preferred: str) -> str:
    """
    지정된 컬럼이 없을 경우, 일반적인 텍스트 컬럼 이름을 순서대로 탐색한다.
    """
    if preferred in columns:
        return preferred
    for candidate in ("text", "content", "document", "body", "passage"):
        if candidate in columns:
            print(
                f"  [INFO] 컬럼 '{preferred}' 미존재 → '{candidate}' 사용. "
                f"(전체 컬럼: {columns[:10]})"
            )
            return candidate
    # 마지막 수단: 첫 번째 문자열 컬럼
    print(
        f"  [경고] 텍스트 컬럼을 찾지 못함. 첫 번째 컬럼 '{columns[0]}' 사용.",
        file=sys.stderr,
    )
    return columns[0]


def get_text_iterator(
    input_dir: Path,
    text_col: str,
    num_proc: int,
) -> tuple[str, Iterator[str]]:
    """
    포맷을 자동 감지하고 알맞은 텍스트 이터레이터를 반환한다.

    Returns:
        (fmt, iterator)  fmt은 감지된 포맷 문자열
    """
    fmt = _detect_format(input_dir)
    if fmt == "hf_arrow":
        return fmt, _iter_hf_arrow(input_dir, text_col, num_proc)
    elif fmt == "parquet":
        return fmt, _iter_parquet(input_dir, text_col)
    elif fmt == "jsonl":
        return fmt, _iter_jsonl(input_dir, text_col)
    else:
        raise RuntimeError(
            f"지원하지 않는 포맷이거나 인식할 수 없습니다: {input_dir}\n"
            f"지원 포맷: HuggingFace arrow, parquet, jsonl"
        )


# ===========================================================================
# 단일 프로세스 토큰화 워커 (multiprocessing.Pool에서 호출)
# ===========================================================================

# 전역 토크나이저 — 각 워커 프로세스에서 한 번만 초기화
_g_sp: SPTokenizer | None = None
_g_model_path: str = ""


def _worker_init(model_path: str) -> None:
    """워커 초기화 함수: SentencePiece 모델 로드."""
    global _g_sp, _g_model_path
    _g_model_path = model_path
    _g_sp = SPTokenizer(model_path)
    _g_sp._ensure_loaded()


def _worker_tokenize_batch(texts: list[str]) -> list[list[int]]:
    """
    텍스트 배치를 토큰화하고 품질 필터를 적용한다.

    반환값: 유효한 토큰 리스트 목록 (필터 통과한 것만)
    """
    global _g_sp
    results: list[list[int]] = []
    for text in texts:
        if not text or not isinstance(text, str):
            continue
        # 품질 필터: 언어
        if not _has_enough_korean_or_english(text):
            continue
        try:
            ids = _g_sp.encode(text)
        except Exception:
            continue
        # 길이 필터
        if len(ids) < MIN_TOKENS:
            continue
        if len(ids) > MAX_TOKENS:
            ids = ids[:MAX_TOKENS]
        results.append(ids)
    return results


# ===========================================================================
# memmap 청크 기반 기록기
# ===========================================================================

class MemmapWriter:
    """
    uint16 numpy memmap 파일에 토큰을 청크 단위로 기록하는 래퍼.

    초기에 작은 크기로 생성하고, 필요할 때 resize한다.
    최종적으로 실제 기록된 크기로 truncate하여 저장한다.
    """

    def __init__(self, path: Path, initial_size: int = CHUNK_TOKENS) -> None:
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        self._alloc = max(initial_size, CHUNK_TOKENS)
        self._mm = np.memmap(
            str(path), dtype="uint16", mode="w+", shape=(self._alloc,)
        )
        self._pos = 0

    def write(self, tokens: Iterable[int]) -> int:
        """tokens를 기록하고 기록된 토큰 수를 반환한다."""
        arr = np.asarray(list(tokens), dtype=np.uint16)
        n = len(arr)
        if n == 0:
            return 0
        needed = self._pos + n
        if needed > self._alloc:
            # 두 배 또는 필요한 크기 중 큰 값으로 확장
            new_alloc = max(self._alloc * 2, needed + CHUNK_TOKENS)
            self._mm.flush()
            del self._mm
            self._alloc = new_alloc
            self._mm = np.memmap(
                str(self.path), dtype="uint16", mode="r+", shape=(self._alloc,)
            )
        self._mm[self._pos : self._pos + n] = arr
        self._pos += n
        return n

    def finalize(self) -> int:
        """기록된 실제 크기로 파일을 truncate하고 닫는다. 총 토큰 수를 반환한다."""
        self._mm.flush()
        del self._mm
        # 실제 기록된 크기로 truncate
        final_bytes = self._pos * 2  # uint16 = 2 bytes
        with open(str(self.path), "r+b") as fh:
            fh.truncate(final_bytes)
        return self._pos


# ===========================================================================
# 핵심 토큰화 파이프라인
# ===========================================================================

def tokenize_directory(
    input_dir: Path,
    output_path: Path,
    tokenizer_path: str,
    text_col: str = "text",
    num_proc: int = 8,
    batch_size: int = 512,
    eos_between_docs: bool = True,
    val_split: float = 0.002,
    seed: int = 42,
) -> dict:
    """
    단일 디렉토리를 토큰화하여 .bin 파일(들)로 저장한다.

    Args:
        input_dir:       입력 디렉토리 (포맷 자동 감지)
        output_path:     출력 .bin 파일 경로 (훈련 셋)
        tokenizer_path:  SentencePiece .model 파일 경로
        text_col:        텍스트 컬럼 이름 (arrow/parquet에서 사용)
        num_proc:        병렬 워커 수
        batch_size:      워커당 배치 크기
        eos_between_docs: 문서 사이에 EOS 토큰 삽입 여부
        val_split:       검증 분리 비율 (0 이면 val 파일 생성 안 함)
        seed:            재현성 시드

    Returns:
        통계 dict (total_tokens, train_tokens, val_tokens, skipped, elapsed_s)
    """
    t_start = time.time()

    # ─── 토크나이저 로드 (메인 프로세스: EOS id 확인) ─────────────────────
    sp_main = SPTokenizer(tokenizer_path)
    eos_id = sp_main.eos_id
    vocab_size = sp_main.vocab_size
    print(f"  토크나이저: {tokenizer_path}")
    print(f"  vocab_size={vocab_size:,}, eos_id={eos_id}")
    if vocab_size > UINT16_MAX:
        print(
            f"  [경고] vocab_size({vocab_size}) > {UINT16_MAX} "
            f"— uint16 오버플로 가능. 65535 이하 id만 안전.",
            file=sys.stderr,
        )

    # ─── 포맷 감지 & 이터레이터 생성 ─────────────────────────────────────
    fmt, text_iter = get_text_iterator(input_dir, text_col, num_proc)
    print(f"  포맷: {fmt}")

    # ─── 출력 경로 설정 ────────────────────────────────────────────────────
    train_path = output_path
    val_path: Path | None = None
    if val_split > 0:
        stem = output_path.stem
        if "train" in stem:
            val_path = output_path.parent / output_path.name.replace("train", "val")
        else:
            val_path = output_path.with_name(stem + "_val" + output_path.suffix)

    print(f"  출력(train): {train_path}")
    if val_path:
        print(f"  출력(val):   {val_path}")

    # ─── memmap 기록기 초기화 ─────────────────────────────────────────────
    writer = MemmapWriter(train_path)
    val_writer: MemmapWriter | None = MemmapWriter(val_path) if val_path else None

    # ─── multiprocessing Pool 생성 ────────────────────────────────────────
    pool = mp.Pool(
        processes=num_proc,
        initializer=_worker_init,
        initargs=(tokenizer_path,),
    )

    total_docs = 0
    skipped    = 0
    total_toks = 0

    # numpy rng for deterministic val split
    rng = np.random.default_rng(seed)

    def _submit_batch(batch_texts: list[str]) -> None:
        nonlocal total_docs, skipped, total_toks
        # 동기 map (배치 단위, 워커별 서브배치로 분할)
        sub_size = max(1, len(batch_texts) // num_proc)
        sub_batches = [
            batch_texts[i : i + sub_size]
            for i in range(0, len(batch_texts), sub_size)
        ]
        results_list = pool.map(_worker_tokenize_batch, sub_batches)

        for results in results_list:
            for ids in results:
                total_docs += 1
                n = len(ids)
                total_toks += n
                # EOS 토큰 삽입
                if eos_between_docs:
                    ids_out = ids + [eos_id]
                else:
                    ids_out = ids

                # val split: 무작위로 val_split 비율만큼 val 파일로
                if val_writer is not None and rng.random() < val_split:
                    val_writer.write(ids_out)
                else:
                    writer.write(ids_out)

            skipped_in_batch = sum(1 for _ in results) - len(results)

    # ─── 배치 수집 & tqdm 진행률 ─────────────────────────────────────────
    batch_buf: list[str] = []
    pbar = tqdm(desc=f"토큰화 [{input_dir.name}]", unit="doc", dynamic_ncols=True)

    for text in text_iter:
        batch_buf.append(text)
        if len(batch_buf) >= batch_size * num_proc:
            _submit_batch(batch_buf)
            pbar.update(len(batch_buf))
            pbar.set_postfix(
                tokens=f"{total_toks:,}",
                docs=f"{total_docs:,}",
                refresh=False,
            )
            batch_buf = []

    # 마지막 잔여 배치 처리
    if batch_buf:
        _submit_batch(batch_buf)
        pbar.update(len(batch_buf))

    pbar.close()
    pool.close()
    pool.join()

    # ─── 파일 마무리 ──────────────────────────────────────────────────────
    train_tokens = writer.finalize()
    val_tokens   = val_writer.finalize() if val_writer else 0

    elapsed = time.time() - t_start
    total_toks_with_eos = train_tokens + val_tokens

    print()
    print(f"  완료: {elapsed:.1f}초")
    print(f"  처리 문서: {total_docs:,}")
    print(f"  총 토큰(EOS 포함): {total_toks_with_eos:,}")
    print(f"    train: {train_tokens:,}  ({train_tokens*2/1e9:.2f} GB)")
    if val_tokens:
        print(f"    val:   {val_tokens:,}  ({val_tokens*2/1e9:.2f} GB)")
    throughput = total_toks_with_eos / elapsed if elapsed > 0 else 0
    print(f"  처리율: {throughput/1e6:.2f} M token/s")

    return {
        "total_docs"   : total_docs,
        "total_tokens" : total_toks_with_eos,
        "train_tokens" : train_tokens,
        "val_tokens"   : val_tokens,
        "elapsed_s"    : elapsed,
        "train_path"   : str(train_path),
        "val_path"     : str(val_path) if val_path else None,
    }


# ===========================================================================
# 서브디렉토리 자동 스캔 모드
# ===========================================================================

def auto_scan_and_tokenize(
    root_dir: Path,
    output_dir: Path,
    tokenizer_path: str,
    text_col: str,
    num_proc: int,
    batch_size: int,
    val_split: float,
    seed: int,
) -> list[dict]:
    """
    root_dir 의 직접 자식 디렉토리를 스캔하여 각각 토큰화한다.

    각 서브디렉토리에 대해:
        output_dir/korean_extra_{subdir_name}_train.bin 을 생성한다.
    """
    children = sorted(p for p in root_dir.iterdir() if p.is_dir())
    if not children:
        raise RuntimeError(f"서브디렉토리가 없습니다: {root_dir}")

    print(f"자동 스캔: {len(children)}개 서브디렉토리 발견")
    for ch in children:
        print(f"  - {ch.name}")
    print()

    all_stats = []
    for child in children:
        print("=" * 60)
        print(f"처리 중: {child}")
        print("=" * 60)
        safe_name = child.name.replace("/", "_").replace(" ", "_")
        out_name  = f"korean_extra_{safe_name}_train.bin"
        out_path  = output_dir / out_name
        try:
            stats = tokenize_directory(
                input_dir      = child,
                output_path    = out_path,
                tokenizer_path = tokenizer_path,
                text_col       = text_col,
                num_proc       = num_proc,
                batch_size     = batch_size,
                val_split      = val_split,
                seed           = seed,
            )
            stats["source"] = child.name
            all_stats.append(stats)
        except Exception as exc:
            print(f"  [오류] {child.name} 처리 실패: {exc}", file=sys.stderr)
            all_stats.append({"source": child.name, "error": str(exc)})
        print()

    return all_stats


# ===========================================================================
# CLI
# ===========================================================================

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "korean_extra/ 대용량 데이터셋을 병렬 토큰화하여 uint16 memmap(.bin) 로 저장. "
            "HuggingFace arrow, parquet, jsonl 포맷 자동 감지."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # 입력
    parser.add_argument(
        "--input_dir",
        required=True,
        help="토큰화할 디렉토리 경로. --auto_scan 시에는 루트 디렉토리.",
    )
    parser.add_argument(
        "--auto_scan",
        action="store_true",
        help=(
            "input_dir 의 직접 자식 디렉토리를 모두 순차 처리. "
            "이 경우 --output_dir 을 지정해야 함."
        ),
    )
    parser.add_argument(
        "--text_col",
        default="text",
        help="텍스트 컬럼 이름 (arrow/parquet/jsonl). 자동 추정 가능.",
    )

    # 출력
    out_group = parser.add_mutually_exclusive_group()
    out_group.add_argument(
        "--output",
        default=None,
        help="출력 .bin 파일 경로 (단일 디렉토리 처리 시 사용).",
    )
    out_group.add_argument(
        "--output_dir",
        default=None,
        help="출력 .bin 파일들을 저장할 디렉토리 (--auto_scan 시 사용).",
    )

    # 토크나이저
    parser.add_argument(
        "--tokenizer",
        default=(
            "/PROJECT/0325120031_A/ghong/taketimes/llm-bang"
            "/tokenizer/korean_64k.model"
        ),
        help="SentencePiece .model 파일 경로.",
    )

    # 처리 옵션
    parser.add_argument(
        "--num_proc",
        type=int,
        default=8,
        help="병렬 워커 수 (multiprocessing.Pool).",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=512,
        help="워커당 배치 크기 (문서 수).",
    )
    parser.add_argument(
        "--val_split",
        type=float,
        default=0.002,
        help="검증 분리 비율 (0.0 이면 val 파일 미생성).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="재현성 시드.",
    )
    parser.add_argument(
        "--no_eos",
        action="store_true",
        help="문서 사이에 EOS 토큰을 삽입하지 않는다.",
    )

    args = parser.parse_args()

    # 검증
    if not args.auto_scan and args.output is None:
        # 자동 출력 경로 생성
        input_name = Path(args.input_dir).name
        args.output = str(
            Path(args.input_dir).parent.parent
            / f"korean_extra_{input_name}_train.bin"
        )
        print(f"[INFO] --output 미지정 → 자동 설정: {args.output}")

    if args.auto_scan and args.output_dir is None:
        parser.error("--auto_scan 사용 시 --output_dir 을 지정해야 합니다.")

    return args


def main() -> None:
    args = parse_args()

    tokenizer_path = args.tokenizer
    if not Path(tokenizer_path).exists():
        # fallback: 상대경로 시도
        fallback = Path(
            "/PROJECT/0325120031_A/ghong/taketimes/llm-bang"
            "/tokenizer/korean_64k.model"
        )
        if fallback.exists():
            tokenizer_path = str(fallback)
        else:
            print(
                f"ERROR: 토크나이저 파일을 찾을 수 없습니다: {tokenizer_path}",
                file=sys.stderr,
            )
            sys.exit(1)

    print("=" * 60)
    print(" LLM-Bang tokenize_extra.py")
    print("=" * 60)
    print(f"  입력:        {args.input_dir}")
    print(f"  토크나이저:  {tokenizer_path}")
    print(f"  num_proc:    {args.num_proc}")
    print(f"  batch_size:  {args.batch_size}")
    print(f"  val_split:   {args.val_split}")
    print(f"  seed:        {args.seed}")
    print(f"  eos:         {not args.no_eos}")
    print()

    if args.auto_scan:
        stats_list = auto_scan_and_tokenize(
            root_dir       = Path(args.input_dir),
            output_dir     = Path(args.output_dir),
            tokenizer_path = tokenizer_path,
            text_col       = args.text_col,
            num_proc       = args.num_proc,
            batch_size     = args.batch_size,
            val_split      = args.val_split,
            seed           = args.seed,
        )
        print("=" * 60)
        print(" 전체 요약")
        print("=" * 60)
        grand_train = 0
        grand_val   = 0
        for s in stats_list:
            if "error" in s:
                print(f"  {s['source']:40s} ERROR: {s['error']}")
            else:
                t = s.get("train_tokens", 0)
                v = s.get("val_tokens", 0)
                grand_train += t
                grand_val   += v
                print(
                    f"  {s['source']:40s} "
                    f"train={t:>14,}  val={v:>12,}  "
                    f"({s['elapsed_s']:.0f}s)"
                )
        print("-" * 60)
        print(
            f"  {'합계':40s} "
            f"train={grand_train:>14,}  val={grand_val:>12,}"
        )
        print(
            f"\n  총 토큰: {grand_train + grand_val:,}  "
            f"({(grand_train + grand_val) * 2 / 1e9:.2f} GB)"
        )

    else:
        stats = tokenize_directory(
            input_dir      = Path(args.input_dir),
            output_path    = Path(args.output),
            tokenizer_path = tokenizer_path,
            text_col       = args.text_col,
            num_proc       = args.num_proc,
            batch_size     = args.batch_size,
            eos_between_docs = not args.no_eos,
            val_split      = args.val_split,
            seed           = args.seed,
        )
        print()
        print("=" * 60)
        print(" 결과 요약")
        print("=" * 60)
        print(f"  train .bin : {stats['train_path']}")
        if stats.get("val_path"):
            print(f"  val .bin   : {stats['val_path']}")
        print(f"  train 토큰 : {stats['train_tokens']:,}")
        print(f"  val 토큰   : {stats['val_tokens']:,}")
        print(f"  처리 문서  : {stats['total_docs']:,}")
        print(f"  소요 시간  : {stats['elapsed_s']:.1f}초")

        # 검증: memmap 로드 테스트
        print()
        print("  [검증] memmap 로드 테스트...")
        try:
            d = np.memmap(stats["train_path"], dtype="uint16", mode="r")
            print(f"  memmap shape: {d.shape}  dtype: {d.dtype}")
            print(f"  첫 10 토큰: {d[:10].tolist()}")
        except Exception as exc:
            print(f"  [경고] memmap 로드 실패: {exc}", file=sys.stderr)


if __name__ == "__main__":
    main()