tps-dashboard/backend/scripts/import_doco_transcripts.py

"""
一次性批量导入脚本：将 doco/deliverables/ 下的 22 期融合A稿 docx 文件
导入知识库（knowledge_items + knowledge_embeddings）。

运行方式：
    cd backend && python -m scripts.import_doco_transcripts
"""

import re
import sys
from pathlib import Path

from docx import Document
from sqlmodel import Session, select

from app.db.session import engine
from app.models.knowledge import KnowledgeItem
from app.services.knowledge_service import KnowledgeService

# ── 路径 ──────────────────────────────────────────────────────
# backend/scripts/ → backend/ → 项目根目录
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
DOCO_DIR = PROJECT_ROOT / "doco" / "deliverables"

# ── 文件名正则 ────────────────────────────────────────────────
# 第{NN}期_{YYYYMMDD}_{节目名}_{编导名}_融合A稿.docx
FILENAME_RE = re.compile(
    r"^第(\d{2,3})期_(\d{8})_(.+?)_(.+?)_融合A稿\.docx$"
)


def parse_filename(filename: str) -> dict | None:
    """从文件名解析期次号、播出日期、节目名、编导名。不匹配返回 None。"""
    m = FILENAME_RE.match(filename)
    if not m:
        return None
    issue_num = int(m.group(1))
    date_raw = m.group(2)  # YYYYMMDD
    title = m.group(3)
    author = m.group(4)
    broadcast_date = f"{date_raw[:4]}-{date_raw[4:6]}-{date_raw[6:8]}"
    return {
        "issue_num": issue_num,
        "broadcast_date": broadcast_date,
        "title": title,
        "author": author,
    }


def extract_docx_text(file_path: Path) -> str:
    """用 python-docx 提取 docx 正文纯文本，段落间用 \\n\\n 分隔。"""
    doc = Document(str(file_path))
    paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
    return "\n\n".join(paragraphs)


def build_md_bytes(title: str, author: str, broadcast_date: str, body: str) -> bytes:
    """组装带 YAML frontmatter 的 md 格式 bytes，供 store_md_file() 使用。"""
    md_content = (
        f"---\n"
        f"类型: 节目文稿\n"
        f"名称: {title}\n"
        f"编导: {author}\n"
        f"播出日期: {broadcast_date}\n"
        f"---\n"
        f"{body}"
    )
    return md_content.encode("utf-8")


def check_exists(source_file_name: str) -> bool:
    """查 knowledge_items 表，判断 source_file_name 是否已存在。"""
    with Session(engine) as session:
        stmt = select(KnowledgeItem).where(
            KnowledgeItem.source_file_name == source_file_name
        )
        existing = session.exec(stmt).first()
        return existing is not None


def main():
    # 确认 docx 目录存在
    if not DOCO_DIR.is_dir():
        print(f"✗ 目录不存在: {DOCO_DIR}")
        sys.exit(1)

    # 收集并排序 docx 文件
    docx_files = sorted(DOCO_DIR.glob("*.docx"), key=lambda p: p.name)
    if not docx_files:
        print("✗ 未找到任何 .docx 文件")
        sys.exit(1)

    print(f"共发现 {len(docx_files)} 个 docx 文件，开始导入...\n")

    service = KnowledgeService()
    success_count = 0
    skip_count = 0
    fail_count = 0
    total = len(docx_files)

    for idx, file_path in enumerate(docx_files, 1):
        filename = file_path.name

        # 解析文件名
        meta = parse_filename(filename)
        if meta is None:
            print(f"[{idx}/{total}] ✗ {filename} — 文件名格式不匹配，跳过")
            fail_count += 1
            continue

        display_label = f"第{meta['issue_num']:02d}期 {meta['title']}"

        # 防重复：用去掉 .docx 后缀加 .md 的文件名
        source_file_name = filename.replace(".docx", ".md")

        try:
            if check_exists(source_file_name):
                print(f"[{idx}/{total}] ⊘ {display_label} — 已存在，跳过")
                skip_count += 1
                continue

            # 提取正文
            body_text = extract_docx_text(file_path)
            char_count = len(body_text)

            # 组装 md bytes
            md_bytes = build_md_bytes(
                title=meta["title"],
                author=meta["author"],
                broadcast_date=meta["broadcast_date"],
                body=body_text,
            )

            # 调用入库链路（含 embedding API 调用）
            service.store_md_file(file_content=md_bytes, file_name=source_file_name)

            print(f"[{idx}/{total}] ✓ {display_label} — 入库成功（{char_count}字）")
            success_count += 1

        except Exception as e:
            print(f"[{idx}/{total}] ✗ {display_label} — 失败：{e}")
            fail_count += 1

    # 汇总
    print(f"\n{'='*50}")
    print(f"导入完成：成功 {success_count} 篇 / 跳过 {skip_count} 篇 / 失败 {fail_count} 篇")
    print(f"{'='*50}")


if __name__ == "__main__":
    main()