wordloop/backend/scripts/import_dictionary_online.py

"""从线上 ECDICT 词库下载并导入 MySQL。

数据源: https://github.com/skywind3000/ECDICT （开源英汉词典，约 340 万词条）

用法（在 backend 目录下）:
  python -m scripts.import_dictionary_online
  python -m scripts.import_dictionary_online --preset full
  python -m scripts.import_dictionary_online --preset standard --skip-download
"""

from __future__ import annotations

import argparse
import sqlite3
import sys
import zipfile
from pathlib import Path
from typing import Optional
from urllib.request import urlretrieve

BACKEND_DIR = Path(__file__).resolve().parent.parent
CACHE_DIR = BACKEND_DIR / "data" / "cache"
ECDICT_ZIP_URL = (
    "https://github.com/skywind3000/ECDICT/releases/download/1.0.28/ecdict-sqlite-28.zip"
)
ECDICT_ZIP_PATH = CACHE_DIR / "ecdict-sqlite-28.zip"
ECDICT_DB_PATH = CACHE_DIR / "stardict.db"

sys.path.insert(0, str(BACKEND_DIR))

from database import SessionLocal  # noqa: E402
from scripts.dictionary_importer import upsert_entries  # noqa: E402

PRESET_SQL = {
    # 常用词：Collins / 牛津核心词，或词频排名前 5 万
    "standard": """
        translation IS NOT NULL AND translation != ''
        AND (
            collins >= 1 OR oxford = 1
            OR (frq IS NOT NULL AND frq <= 50000)
        )
    """,
    # 核心词：Collins 星级或牛津 3000
    "core": """
        translation IS NOT NULL AND translation != ''
        AND (collins >= 1 OR oxford = 1)
    """,
    # 全量（约 338 万，耗时较长）
    "full": """
        translation IS NOT NULL AND translation != ''
    """,
}


def download(url: str, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"正在下载 ECDICT SQLite（约 207MB）…")
    print(f"  {url}")

    def progress(block_num: int, block_size: int, total_size: int) -> None:
        if total_size <= 0:
            return
        done = block_num * block_size
        pct = min(100, done * 100 // total_size)
        mb = done / 1024 / 1024
        total_mb = total_size / 1024 / 1024
        print(f"\r  进度: {pct:3d}% ({mb:.1f}/{total_mb:.1f} MB)", end="", flush=True)

    urlretrieve(url, dest, reporthook=progress)
    print()


def ensure_ecdict_db(skip_download: bool) -> Path:
    if ECDICT_DB_PATH.exists():
        return ECDICT_DB_PATH

    if not ECDICT_ZIP_PATH.exists():
        if skip_download:
            print(f"错误: 未找到本地词库，请先下载或去掉 --skip-download", file=sys.stderr)
            sys.exit(1)
        download(ECDICT_ZIP_URL, ECDICT_ZIP_PATH)

    print(f"正在解压 {ECDICT_ZIP_PATH.name} …")
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(ECDICT_ZIP_PATH, "r") as zf:
        zf.extract("stardict.db", CACHE_DIR)
    print(f"词库就绪: {ECDICT_DB_PATH}")
    return ECDICT_DB_PATH


def count_entries(db_path: Path, where: str) -> int:
    conn = sqlite3.connect(db_path)
    try:
        cur = conn.execute(f"SELECT COUNT(*) FROM stardict WHERE {where}")
        return int(cur.fetchone()[0])
    finally:
        conn.close()


def iter_ecdict_rows(db_path: Path, where: str, limit: Optional[int]):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    try:
        sql = f"""
            SELECT word, phonetic, translation
            FROM stardict
            WHERE {where}
            ORDER BY COALESCE(frq, 9999999), word
        """
        if limit:
            sql += f" LIMIT {int(limit)}"
        cur = conn.execute(sql)
        while True:
            rows = cur.fetchmany(5000)
            if not rows:
                break
            for row in rows:
                yield {
                    "word": row["word"],
                    "phonetic": row["phonetic"],
                    "translation": row["translation"],
                }
    finally:
        conn.close()


def import_from_ecdict(
    db_path: Path,
    preset: str,
    limit: Optional[int],
    batch_size: int,
) -> int:
    where = PRESET_SQL[preset]
    total_expected = count_entries(db_path, where)
    if limit:
        total_expected = min(total_expected, limit)
    print(f"预设: {preset}，预计导入 {total_expected:,} 条")

    session = SessionLocal()
    imported = 0
    buffer: list[dict] = []
    try:
        for row in iter_ecdict_rows(db_path, where, limit):
            buffer.append(row)
            if len(buffer) >= batch_size:
                imported += upsert_entries(session, buffer, source="ecdict", batch_size=batch_size)
                buffer.clear()
                print(f"\r  已导入 {imported:,} / {total_expected:,}", end="", flush=True)
        if buffer:
            imported += upsert_entries(session, buffer, source="ecdict", batch_size=batch_size)
        print()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()
    return imported


def main() -> None:
    parser = argparse.ArgumentParser(description="从线上 ECDICT 下载并导入离线翻译词典")
    parser.add_argument(
        "--preset",
        choices=list(PRESET_SQL.keys()),
        default="standard",
        help="standard=常用约 83 万 | core=核心约 1.4 万 | full=全量约 338 万（默认 standard）",
    )
    parser.add_argument("--limit", type=int, default=None, help="最多导入条数（调试用）")
    parser.add_argument("--batch-size", type=int, default=2000, help="MySQL 批量写入大小")
    parser.add_argument("--skip-download", action="store_true", help="跳过下载，使用本地 cache")
    args = parser.parse_args()

    db_path = ensure_ecdict_db(args.skip_download)
    imported = import_from_ecdict(db_path, args.preset, args.limit, args.batch_size)
    print(f"导入完成: {imported:,} 条（source=ecdict, preset={args.preset}）")


if __name__ == "__main__":
    main()