#!/usr/bin/env bash
# =============================================================================
# prepare_3b_data.sh — 3B 모델 학습 데이터 전체 파이프라인
#
# 사용법:
#   bash scripts/prepare_3b_data.sh [--step N] [--jobs 72]
#
# 스텝:
#   1 = CulturaX 토큰화
#   2 = cc100 해제 + 토큰화
#   3 = OSCAR 토큰화
#   4 = korean_webtext 토큰화
#   5 = HPLT 한국어 추출 + 토큰화
#   6 = textbooks + finepdfs + kovast 토큰화
#   7 = 전체 병합
#   8 = train/val split 검증
# =============================================================================
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
cd "${PROJECT_ROOT}"

# ─── 설정 ────────────────────────────────────────────────────────────────
DATA_DIR="data"
EXTRA_DIR="data/korean_extra"
TOKENIZER="tokenizer/tokenizer.json"
VAL_SPLIT=0.002
SEED=42
JOBS=72
FROM_STEP=0
LOG_FILE="data/prepare_3b.log"

while [[ $# -gt 0 ]]; do
    case $1 in
        --step)   FROM_STEP="$2"; shift 2 ;;
        --jobs)   JOBS="$2"; shift 2 ;;
        *)        echo "Unknown arg: $1"; exit 1 ;;
    esac
done

mkdir -p "$(dirname "$LOG_FILE")"
exec > >(tee -a "$LOG_FILE") 2>&1

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

# ─── 토큰화 헬퍼 (parquet → bin) ─────────────────────────────────────────
tokenize_parquet() {
    local name="$1"
    local input_pattern="$2"
    local text_col="$3"
    local output="${DATA_DIR}/${name}_train.bin"

    if [[ -f "$output" && $FROM_STEP -le 0 ]]; then
        log "[SKIP] $output already exists ($(du -h "$output" | cut -f1))"
        return
    fi

    log "[START] Tokenizing $name from parquet..."
    python3 - <<PYEOF
import glob, os, sys
import numpy as np
from tokenizers import Tokenizer
import pyarrow.parquet as pq
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor
import multiprocessing as mp

tokenizer_path = "${TOKENIZER}"
input_pattern = "${input_pattern}"
text_col = "${text_col}"
output_train = "${output}"
output_val = output_train.replace("_train.bin", "_val.bin")
val_split = ${VAL_SPLIT}
seed = ${SEED}

files = sorted(glob.glob(input_pattern))
print(f"Found {len(files)} parquet files")

tokenizer = Tokenizer.from_file(tokenizer_path)

all_tokens = []
total_docs = 0

for f in tqdm(files, desc="${name}"):
    try:
        table = pq.read_table(f, columns=[text_col])
        for text in table.column(text_col):
            t = text.as_py()
            if t and len(t) > 50:
                ids = tokenizer.encode(t).ids
                all_tokens.extend(ids)
                total_docs += 1
    except Exception as e:
        print(f"Error processing {f}: {e}", file=sys.stderr)
        continue

print(f"Total: {total_docs:,} docs, {len(all_tokens):,} tokens")

# Split
import random
random.seed(seed)
random.shuffle(all_tokens)  # Not ideal but matches existing code
n_val = int(len(all_tokens) * val_split)
val_tokens = all_tokens[:n_val]
train_tokens = all_tokens[n_val:]

np.array(train_tokens, dtype=np.uint16).tofile(output_train)
np.array(val_tokens, dtype=np.uint16).tofile(output_val)
print(f"Saved: {output_train} ({len(train_tokens):,} tokens)")
print(f"Saved: {output_val} ({len(val_tokens):,} tokens)")
PYEOF
    log "[DONE] $name → $output"
}

# ─── Step 1: CulturaX ────────────────────────────────────────────────────
if [[ $FROM_STEP -le 1 ]]; then
    log "=== Step 1: CulturaX 토큰화 ==="
    tokenize_parquet "culturax" \
        "${EXTRA_DIR}/culturax_ko/ko/*.parquet" \
        "text"
fi

# ─── Step 2: cc100 해제 + 토큰화 ─────────────────────────────────────────
if [[ $FROM_STEP -le 2 ]]; then
    log "=== Step 2: cc100 해제 + 토큰화 ==="
    CC100_XZ="${EXTRA_DIR}/cc100_ko/ko.txt.xz"
    CC100_TXT="${EXTRA_DIR}/cc100_ko/ko.txt"
    CC100_OUT="${DATA_DIR}/cc100_train.bin"

    if [[ -f "$CC100_OUT" && $FROM_STEP -le 0 ]]; then
        log "[SKIP] cc100 already tokenized"
    else
        # 해제
        if [[ ! -f "$CC100_TXT" ]]; then
            log "Decompressing cc100 xz (14GB → 54GB)..."
            xz -dk "$CC100_XZ"
            log "Decompression done"
        fi

        # 토큰화 (대용량 → 스트리밍)
        log "Tokenizing cc100 (54GB text)..."
        python3 - <<'PYEOF'
import numpy as np
from tokenizers import Tokenizer
from tqdm import tqdm
import random

tokenizer = Tokenizer.from_file("tokenizer/tokenizer.json")
input_file = "data/korean_extra/cc100_ko/ko.txt"
output_train = "data/cc100_train.bin"
output_val = "data/cc100_val.bin"

# Stream tokenize in chunks
all_tokens = []
doc_buffer = []
doc_count = 0

with open(input_file, 'r', encoding='utf-8', errors='replace') as f:
    for line in tqdm(f, desc="cc100", unit=" lines"):
        line = line.strip()
        if not line:
            # Document boundary
            if doc_buffer:
                text = '\n'.join(doc_buffer)
                if len(text) > 50:
                    ids = tokenizer.encode(text).ids
                    all_tokens.extend(ids)
                    doc_count += 1
                doc_buffer = []
        else:
            doc_buffer.append(line)

    # Last doc
    if doc_buffer:
        text = '\n'.join(doc_buffer)
        if len(text) > 50:
            all_tokens.extend(tokenizer.encode(text).ids)
            doc_count += 1

print(f"Total: {doc_count:,} docs, {len(all_tokens):,} tokens")

# Split
n_val = int(len(all_tokens) * 0.002)
np.array(all_tokens[n_val:], dtype=np.uint16).tofile(output_train)
np.array(all_tokens[:n_val], dtype=np.uint16).tofile(output_val)
print(f"Saved train: {len(all_tokens)-n_val:,} tokens")
print(f"Saved val: {n_val:,} tokens")
PYEOF
        log "[DONE] cc100"
    fi
fi

# ─── Step 3: OSCAR ───────────────────────────────────────────────────────
if [[ $FROM_STEP -le 3 ]]; then
    log "=== Step 3: OSCAR 토큰화 ==="
    OSCAR_OUT="${DATA_DIR}/oscar_train.bin"

    if [[ -f "$OSCAR_OUT" && $FROM_STEP -le 0 ]]; then
        log "[SKIP] OSCAR already tokenized"
    else
        python3 - <<'PYEOF'
import glob, numpy as np
from tokenizers import Tokenizer
import pyarrow.parquet as pq
from tqdm import tqdm

tokenizer = Tokenizer.from_file("tokenizer/tokenizer.json")
files = sorted(glob.glob("data/korean_extra/oscar_ko/data/kor_Hang/*.parquet"))
all_tokens = []
doc_count = 0

for f in tqdm(files, desc="OSCAR"):
    table = pq.read_table(f, columns=['text'])
    for row in table.column('text'):
        if row is None:
            continue
        parts = row.as_py()
        if parts:
            text = '\n'.join(item['text'] for item in parts if item and item.get('text'))
            if len(text) > 50:
                all_tokens.extend(tokenizer.encode(text).ids)
                doc_count += 1

print(f"OSCAR: {doc_count:,} docs, {len(all_tokens):,} tokens")
n_val = int(len(all_tokens) * 0.002)
np.array(all_tokens[n_val:], dtype=np.uint16).tofile("data/oscar_train.bin")
np.array(all_tokens[:n_val], dtype=np.uint16).tofile("data/oscar_val.bin")
PYEOF
        log "[DONE] OSCAR"
    fi
fi

# ─── Step 4: korean_webtext ──────────────────────────────────────────────
if [[ $FROM_STEP -le 4 ]]; then
    log "=== Step 4: korean_webtext 토큰화 ==="
    tokenize_parquet "webtext" \
        "${EXTRA_DIR}/korean_webtext/data/*.parquet" \
        "text"
fi

# ─── Step 5: HPLT 한국어 추출 + 토큰화 ──────────────────────────────────
if [[ $FROM_STEP -le 5 ]]; then
    log "=== Step 5: HPLT 한국어 추출 + 토큰화 ==="
    HPLT_OUT="${DATA_DIR}/hplt_ko_train.bin"

    if [[ -f "$HPLT_OUT" && $FROM_STEP -le 0 ]]; then
        log "[SKIP] HPLT already tokenized"
    else
        python3 - <<'PYEOF'
import glob, numpy as np
from tokenizers import Tokenizer
import pyarrow.parquet as pq
from tqdm import tqdm

tokenizer = Tokenizer.from_file("tokenizer/tokenizer.json")
files = sorted(glob.glob("data/korean_extra/hplt_ko/en-ko/*.parquet"))
all_tokens = []
doc_count = 0

for f in tqdm(files, desc="HPLT"):
    table = pq.read_table(f, columns=['tgt_doc'])
    for row in table.column('tgt_doc'):
        d = row.as_py()
        if d and d.get('sentences'):
            text = '\n'.join(s for s in d['sentences'] if s)
            if len(text) > 50:
                all_tokens.extend(tokenizer.encode(text).ids)
                doc_count += 1

print(f"HPLT Korean: {doc_count:,} docs, {len(all_tokens):,} tokens")
n_val = int(len(all_tokens) * 0.002)
np.array(all_tokens[n_val:], dtype=np.uint16).tofile("data/hplt_ko_train.bin")
np.array(all_tokens[:n_val], dtype=np.uint16).tofile("data/hplt_ko_val.bin")
PYEOF
        log "[DONE] HPLT"
    fi
fi

# ─── Step 6: textbooks + finepdfs + kovast ───────────────────────────────
if [[ $FROM_STEP -le 6 ]]; then
    log "=== Step 6: 기타 소스 토큰화 ==="
    EXTRA_OUT="${DATA_DIR}/extra_misc_train.bin"

    if [[ -f "$EXTRA_OUT" && $FROM_STEP -le 0 ]]; then
        log "[SKIP] extra_misc already tokenized"
    else
        python3 - <<'PYEOF'
import glob, numpy as np, os
from tokenizers import Tokenizer
import pyarrow.parquet as pq
from tqdm import tqdm

tokenizer = Tokenizer.from_file("tokenizer/tokenizer.json")
all_tokens = []
doc_count = 0

# korean_textbooks (MMLU-style: look for text columns)
tb_files = glob.glob("data/korean_extra/korean_textbooks/**/*.parquet", recursive=True)
for f in tqdm(tb_files, desc="textbooks"):
    try:
        table = pq.read_table(f)
        # Try common text columns
        for col in ['question', 'text', 'input', 'instruction']:
            if col in table.column_names:
                for val in table.column(col):
                    t = val.as_py()
                    if t and len(t) > 20:
                        all_tokens.extend(tokenizer.encode(t).ids)
                        doc_count += 1
                break
    except:
        continue

# finepdfs
pdf_files = glob.glob("data/korean_extra/finepdfs_edu_ko/*.parquet")
for f in tqdm(pdf_files, desc="finepdfs"):
    try:
        table = pq.read_table(f)
        for col in ['text', 'content']:
            if col in table.column_names:
                for val in table.column(col):
                    t = val.as_py()
                    if t and len(t) > 50:
                        all_tokens.extend(tokenizer.encode(t).ids)
                        doc_count += 1
                break
    except:
        continue

print(f"Extra: {doc_count:,} docs, {len(all_tokens):,} tokens")
n_val = int(len(all_tokens) * 0.002)
np.array(all_tokens[n_val:], dtype=np.uint16).tofile("data/extra_misc_train.bin")
np.array(all_tokens[:n_val], dtype=np.uint16).tofile("data/extra_misc_val.bin")
PYEOF
        log "[DONE] extra_misc"
    fi
fi

# ─── Step 7: 전체 병합 ──────────────────────────────────────────────────
if [[ $FROM_STEP -le 7 ]]; then
    log "=== Step 7: 전체 병합 ==="

    TRAIN_BINS=""
    for f in \
        "${DATA_DIR}/korean_train.bin" \
        "${DATA_DIR}/culturax_train.bin" \
        "${DATA_DIR}/cc100_train.bin" \
        "${DATA_DIR}/oscar_train.bin" \
        "${DATA_DIR}/webtext_train.bin" \
        "${DATA_DIR}/hplt_ko_train.bin" \
        "${DATA_DIR}/extra_misc_train.bin"; do
        if [[ -f "$f" ]]; then
            TRAIN_BINS="$TRAIN_BINS $f"
            log "  Including: $f ($(du -h "$f" | cut -f1))"
        else
            log "  [WARN] Missing: $f"
        fi
    done

    if [[ -n "$TRAIN_BINS" ]]; then
        python3 data/merge_bins.py $TRAIN_BINS "${DATA_DIR}/merged_3b_train.bin"
        log "[DONE] merged_3b_train.bin created"
    fi

    # Val 병합
    VAL_BINS=""
    for f in \
        "${DATA_DIR}/korean_val.bin" \
        "${DATA_DIR}/culturax_val.bin" \
        "${DATA_DIR}/cc100_val.bin" \
        "${DATA_DIR}/oscar_val.bin" \
        "${DATA_DIR}/webtext_val.bin" \
        "${DATA_DIR}/hplt_ko_val.bin" \
        "${DATA_DIR}/extra_misc_val.bin"; do
        if [[ -f "$f" ]]; then
            VAL_BINS="$VAL_BINS $f"
        fi
    done

    if [[ -n "$VAL_BINS" ]]; then
        python3 data/merge_bins.py $VAL_BINS "${DATA_DIR}/merged_3b_val.bin"
        log "[DONE] merged_3b_val.bin created"
    fi
fi

# ─── Step 8: 검증 ────────────────────────────────────────────────────────
if [[ $FROM_STEP -le 8 ]]; then
    log "=== Step 8: 최종 검증 ==="
    python3 - <<'PYEOF'
import os, glob
import numpy as np

print("=== 토큰화 결과 ===")
total_train = 0
total_val = 0
for f in sorted(glob.glob("data/*_train.bin") + glob.glob("data/train.bin")):
    n = os.path.getsize(f) // 2
    total_train += n
    print(f"  {os.path.basename(f):30s}: {n:>15,} tokens ({os.path.getsize(f)/1e9:.2f} GB)")

for f in sorted(glob.glob("data/*_val.bin") + glob.glob("data/val.bin")):
    n = os.path.getsize(f) // 2
    total_val += n

print(f"\n  Total train: {total_train:,} tokens ({total_train/1e9:.1f}B)")
print(f"  Total val:   {total_val:,} tokens ({total_val/1e6:.1f}M)")
print(f"\n  3B Chinchilla minimum: 60B tokens")
print(f"  Epochs needed for 60B: {60e9/total_train:.1f}")
print(f"  Epochs needed for 100B: {100e9/total_train:.1f}")
PYEOF
fi

log "=== 파이프라인 완료 ==="