wordloop/backend/scripts/seed_user_words.py

"""向用户词库批量导入单词（测试用）。

用法（backend 目录）:
  python -m scripts.seed_user_words --username john --count 200
  python -m scripts.seed_user_words --username john --count 200 --with-quiz 80
"""

from __future__ import annotations

import argparse
import random
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

BACKEND_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BACKEND_DIR))

from database import SessionLocal  # noqa: E402
from models import DictionaryEntry, QuizRecord, User, Word  # noqa: E402
from services.word_service import calc_mastery_score, utc_now_iso  # noqa: E402


def utc_now_iso_at(dt: datetime) -> str:
    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")


def seed_words(
    db,
    user: User,
    count: int,
    offset: int,
) -> tuple[int, int]:
    existing_pairs = {
        (w.source_text, w.target_text)
        for w in db.query(Word).filter(Word.user_id == user.id).all()
    }
    entries = (
        db.query(DictionaryEntry)
        .order_by(DictionaryEntry.id)
        .offset(offset)
        .limit(count * 3)
        .all()
    )
    added = 0
    skipped = 0
    now = utc_now_iso()

    for e in entries:
        if added >= count:
            break
        zh = (e.zh or "").strip()
        en = (e.lemma_en or "").strip().lower()
        if not zh or not en:
            skipped += 1
            continue
        if (en, zh) in existing_pairs or (zh, en) in existing_pairs:
            skipped += 1
            continue

        word = Word(
            user_id=user.id,
            source_text=zh,
            target_text=en,
            source_lang="zh",
            target_lang="en",
            phonetic=e.phonetic,
            example_en=e.example_en,
            example_cn=e.example_cn,
            status=random.choice(["new", "new", "learning", "mastered", "weak"]),
            correct_count=random.randint(0, 8),
            wrong_count=random.randint(0, 4),
            consecutive_correct_count=random.randint(0, 3),
            mastery_score=0,
            train_count=0,
            total_train_seconds=0,
            review_due_date=datetime.now(timezone.utc).strftime("%Y-%m-%d"),
            created_at=now,
        )
        word.mastery_score = calc_mastery_score(word.correct_count, word.wrong_count)
        db.add(word)
        existing_pairs.add((zh, en))
        added += 1

    db.commit()
    return added, skipped


def seed_quiz_records(db, user: User, max_records: int) -> int:
    words = db.query(Word).filter(Word.user_id == user.id).all()
    if not words:
        return 0
    now = datetime.now(timezone.utc)
    created = 0
    for i in range(max_records):
        w = random.choice(words)
        is_correct = random.random() > 0.35
        t = now - timedelta(days=random.randint(0, 20), hours=random.randint(0, 23))
        record = QuizRecord(
            user_id=user.id,
            word_id=w.id,
            question_type=random.choice(["en_to_zh", "zh_to_en", "spell"]),
            user_answer="test",
            correct_answer="test",
            is_correct=1 if is_correct else 0,
            duration_seconds=random.randint(3, 45),
            created_at=utc_now_iso_at(t),
        )
        db.add(record)
        created += 1
    db.commit()
    return created


def main() -> None:
    parser = argparse.ArgumentParser(description="批量导入用户单词")
    parser.add_argument("--username", default="john")
    parser.add_argument("--count", type=int, default=200)
    parser.add_argument("--offset", type=int, default=5000, help="词典偏移，避免总取相同词")
    parser.add_argument("--with-quiz", type=int, default=0, help="额外生成 N 条模拟训练记录")
    args = parser.parse_args()

    db = SessionLocal()
    try:
        user = db.query(User).filter(User.username == args.username).first()
        if not user:
            print(f"用户不存在: {args.username}", file=sys.stderr)
            sys.exit(1)

        before = db.query(Word).filter(Word.user_id == user.id).count()
        added, skipped = seed_words(db, user, args.count, args.offset)
        after = db.query(Word).filter(Word.user_id == user.id).count()
        print(f"用户 {args.username}: 原有 {before} 词, 新增 {added}, 跳过 {skipped}, 现有 {after} 词")

        if args.with_quiz > 0:
            n = seed_quiz_records(db, user, args.with_quiz)
            print(f"已生成模拟训练记录 {n} 条")
    finally:
        db.close()


if __name__ == "__main__":
    main()