#!/usr/bin/env python3
"""下载维基百科中文数据 + 分词 + 生成 LAL 训练 .bin 文件.

用法:
  python3 scripts/prepare_wiki_data.py [--n_articles 10000] [--out data/wiki_bpe.bin]

输出格式: LALT 二进制 (与 large_bpe_v3.bin 兼容)
  [magic: 4 bytes = "LALT"]
  [n_samples: int32]
  [n_vocab: int32]
  then n_samples records:
    [n_tokens: int32]
    [token_ids: int32 * n_tokens]
"""
import os
import sys
import struct
import subprocess
import argparse
import time

# === 固化配置 ===
TOKENIZER_MODEL = "tokenizer/chinese_bpe.model"
WIKI_DUMP_URL = "https://dumps.wikimedia.org/zhwiki/latest/zhwiki-latest-pages-articles-multistream.xml.bz2"
DEFAULT_OUT = "data/wiki_bpe.bin"
DEFAULT_N_ARTICLES = 10000  # 先取 1 万篇, 约 500-1000 万 token


def log(msg):
    print(f"[WIKI] {msg}", flush=True)


def download_wiki_dump(local_path, max_articles=None):
    """下载维基百科 dump (流式解压 + 解析, 避免下载整个 2GB+ bz2)."""
    import bz2
    import xml.etree.ElementTree as ET
    import urllib.request
    
    log(f"下载+解析维基百科 (最多 {max_articles or '全部'} 篇)...")
    
    articles = []
    in_text = False
    in_title = False
    current_title = ""
    current_text = ""
    article_count = 0
    
    # 流式下载 + bz2 解压
    req = urllib.request.Request(WIKI_DUMP_URL, headers={"User-Agent": "LAL-Data-Prep/1.0"})
    with urllib.request.urlopen(req) as resp:
        with bz2.open(resp, "rt", encoding="utf-8") as f:
            for line in f:
                if "<title>" in line:
                    start = line.index("<title>") + 7
                    end = line.index("</title>")
                    current_title = line[start:end].strip()
                elif "<text" in line:
                    in_text = True
                    # 提取 text 标签内的内容
                    if ">" in line:
                        start = line.index(">") + 1
                        current_text = line[start:]
                    else:
                        current_text = ""
                elif "</text>" in line:
                    in_text = False
                    end = line.index("</text>")
                    current_text += line[:end]
                    
                    # 过滤: 跳过重定向、空页面
                    if current_text and not current_text.startswith("#REDIRECT"):
                        # 清理 wiki 标记 (简单版)
                        text = clean_wiki_text(current_text)
                        if len(text) > 100:  # 太短的文章跳过
                            articles.append(text)
                            article_count += 1
                            if article_count % 1000 == 0:
                                log(f"  已收集 {article_count} 篇文章")
                    
                    current_text = ""
                    
                    if max_articles and article_count >= max_articles:
                        break
            
            # 处理最后一篇
            if in_text and current_text:
                text = clean_wiki_text(current_text)
                if len(text) > 100:
                    articles.append(text)
                    article_count += 1
    
    log(f"共收集 {len(articles)} 篇文章")
    return articles


def clean_wiki_text(text):
    """简单清理 wiki 标记."""
    import re
    # 去掉 wiki 模板 {{...}}
    text = re.sub(r'\{\{[^}]*\}\}', '', text)
    # 去掉 wiki 链接 [[...]]
    text = re.sub(r'\[\[([^|\]]*\|)?([^\]]*)\]\]', r'\2', text)
    # 去掉 HTML 标签
    text = re.sub(r'<[^>]+>', '', text)
    # 去掉 wiki 标题标记 ==
    text = re.sub(r'^=+\s*([^=]+)\s*=+$', r'\1', text, flags=re.MULTILINE)
    # 去掉引用
    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
    text = re.sub(r'<ref[^>]*/>', '', text)
    # 去掉多余空行
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


def tokenize_with_bpe(texts, tokenizer_model):
    """用 sentencepiece BPE 分词."""
    try:
        import sentencepiece as spm
    except ImportError:
        log("安装 sentencepiece...")
        subprocess.run([sys.executable, "-m", "pip", "install", "-q", "sentencepiece"], check=True)
        import sentencepiece as spm
    
    sp = spm.SentencePieceProcessor()
    sp.Load(tokenizer_model)
    
    all_samples = []
    total_tokens = 0
    
    for i, text in enumerate(texts):
        # 分词 (每篇文章作为一个 sample)
        tokens = sp.EncodeAsIds(text)
        if len(tokens) > 10:  # 太短的跳过
            all_samples.append(tokens)
            total_tokens += len(tokens)
        
        if (i + 1) % 1000 == 0:
            log(f"  分词 {i+1}/{len(texts)} 篇, 总 token {total_tokens}")
    
    log(f"分词完成: {len(all_samples)} samples, {total_tokens} tokens ({total_tokens/10000:.1f}万)")
    return all_samples


def write_lalt_bin(samples, out_path, n_vocab=32768):
    """写 LALT 二进制格式."""
    log(f"写入 {out_path}...")
    with open(out_path, "wb") as f:
        # header
        f.write(b"LALT")
        f.write(struct.pack("<i", len(samples)))
        f.write(struct.pack("<i", n_vocab))
        
        # samples
        for tokens in samples:
            f.write(struct.pack("<i", len(tokens)))
            for tok in tokens:
                f.write(struct.pack("<i", tok))
    
    size = os.path.getsize(out_path)
    log(f"完成: {out_path} ({size / 1024 / 1024:.1f} MB)")


def main():
    parser = argparse.ArgumentParser(description="准备维基百科训练数据")
    parser.add_argument("--n_articles", type=int, default=DEFAULT_N_ARTICLES,
                        help=f"文章数 (默认 {DEFAULT_N_ARTICLES})")
    parser.add_argument("--out", type=str, default=DEFAULT_OUT,
                        help=f"输出路径 (默认 {DEFAULT_OUT})")
    args = parser.parse_args()
    
    # 检查 tokenizer
    if not os.path.exists(TOKENIZER_MODEL):
        log(f"[!] tokenizer 不存在: {TOKENIZER_MODEL}")
        sys.exit(1)
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    
    start = time.time()
    
    # 1. 下载 + 解析维基百科
    articles = download_wiki_dump(args.out + ".tmp", max_articles=args.n_articles)
    
    # 2. 分词
    samples = tokenize_with_bpe(articles, TOKENIZER_MODEL)
    
    # 3. 写 .bin
    write_lalt_bin(samples, args.out)
    
    elapsed = time.time() - start
    log(f"总计耗时 {elapsed:.0f}s")
    log(f"数据文件: {args.out}")
    log(f"替换训练数据: cp {args.out} data/large_bpe_v3.bin")


if __name__ == "__main__":
    main()