tiku/tiku-server/app/routers/import_.py

"""导入路由"""

import json
import queue
import secrets
import threading
from pathlib import Path

from fastapi import APIRouter, Request, HTTPException, UploadFile, File, Form, Depends, Query
from fastapi.responses import FileResponse

from app.database import (
    get_connection, create_import_task, update_import_task, get_import_task,
    upsert_problem, update_fts_body,
)
from app.routers.auth import verify_token
from app.services.docx_parser import split_docx_to_questions, create_single_docx, extract_plain_text
from app.services.knowledge import load_knowledge_map
from app.services.backup import auto_backup

router = APIRouter()

# 单线程导入队列
_import_queue = queue.Queue()

def _import_worker():
    """单线程消费导入任务，避免并发线程过多导致崩溃"""
    while True:
        task = _import_queue.get()
        if task is None:
            break
        task_id, tmp_path, import_type, source_name, config = task
        try:
            _process_import(task_id, tmp_path, import_type, source_name, config)
        except Exception as e:
            conn = get_connection(config.db_path)
            try:
                update_import_task(conn, task_id, status="error", message=str(e))
            except Exception:
                pass
            finally:
                conn.close()
        _import_queue.task_done()

# 启动后台工作线程
_worker = threading.Thread(target=_import_worker, daemon=True)
_worker.start()


@router.post("/import")
async def import_docx(
    request: Request,
    file: UploadFile = File(...),
    import_type: str = Form("精选题"),
    _=Depends(verify_token),
):
    """上传整卷 docx 进行导入（队列处理）"""
    config = request.app.state.config

    if not file.filename.endswith('.docx'):
        raise HTTPException(status_code=400, detail={
            "error": {"code": "invalid_file", "message": "仅支持 .docx 文件"}
        })

    task_id = secrets.token_hex(8)

    upload_dir = config.data_dir / "uploads"
    upload_dir.mkdir(parents=True, exist_ok=True)
    tmp_path = upload_dir / f"{task_id}.docx"

    content = await file.read()
    with open(tmp_path, "wb") as f:
        f.write(content)

    conn = get_connection(config.db_path)
    try:
        create_import_task(conn, task_id)
    finally:
        conn.close()

    # 放入队列，单线程顺序处理
    _import_queue.put((task_id, tmp_path, import_type, file.filename, config))

    return {"task_id": task_id, "status": "pending"}


def _process_import(task_id: str, tmp_path: Path, import_type: str, source_name: str, config):
    """后台处理导入"""
    from app.database import _db_lock
    with _db_lock:
        conn = get_connection(config.db_path)
        try:
            _do_import(conn, task_id, tmp_path, import_type, source_name, config)
        finally:
            conn.close()


def _do_import(conn, task_id: str, tmp_path: Path, import_type: str, source_name: str, config):
    """执行导入（调用方已持有数据库锁）"""
    try:
        update_import_task(conn, task_id, status="processing", message="正在拆分题目...")

        # 加载知识点映射
        knowledge_map = load_knowledge_map(config)

        # 拆分
        questions = split_docx_to_questions(tmp_path)
        total = len(questions)
        update_import_task(conn, task_id, total=total, message=f"共 {total} 题，正在处理...")

        imported = []
        skipped = []
        errors = []
        imported_files = []  # 记录导入的 docx 文件路径，用于备份

        for i, q in enumerate(questions):
            update_import_task(conn, task_id, completed=i + 1, message=f"处理第 {i + 1}/{total} 题...")

            try:
                # 缺失元数据则使用默认值
                if q["metadata"]["difficulty"] is None:
                    q["metadata"]["difficulty"] = 0.5

                docx_path, pid, info = create_single_docx(tmp_path, q, import_type, source_name, knowledge_map, config.problems_dir)

                # 入库
                conn2 = get_connection(config.db_path)
                try:
                    # 检查是否已存在（去重）
                    existing = conn2.execute("SELECT id FROM problems WHERE id = ?", (pid,)).fetchone()
                    if existing:
                        skipped.append({"num": q["question_num"], "id": pid, "reason": "内容重复"})
                    else:
                        problem = {
                            "id": pid,
                            "file_path": str(docx_path),
                            "mtime": docx_path.stat().st_mtime,
                            "checksum": info["checksum"],
                            "chapters": info["chapters"],
                            "difficulty": info["difficulty"],
                            "question_type": info["question_type"],
                            "import_type": info["import_type"],
                            "knowledge_codes": info["knowledge_codes"],
                            "comment": info["comment"],
                            "source": info["source"],
                        }
                        upsert_problem(conn2, problem)
                        body = extract_plain_text(docx_path)
                        update_fts_body(conn2, pid, body)
                        imported.append({"num": q["question_num"], "id": pid, "type": info["question_type"]})
                        imported_files.append(str(docx_path))
                finally:
                    conn2.close()

            except Exception as e:
                errors.append({"num": q["question_num"], "error": str(e)})

        # 完成
        result = json.dumps({
            "imported": len(imported),
            "skipped": len(skipped),
            "errors": len(errors),
            "details": {
                "imported": imported,
                "skipped": skipped,
                "errors": errors,
            }
        }, ensure_ascii=False)

        update_import_task(conn, task_id, status="done", message="导入完成", result=result)

        # 自动备份（bank.db + 导入的 docx）
        try:
            auto_backup(config, changed_files=imported_files if imported_files else None)
        except Exception:
            pass

        # 清理临时文件
        try:
            tmp_path.unlink()
        except Exception:
            pass

    except Exception as e:
        update_import_task(conn, task_id, status="error", message=str(e))
    finally:
        conn.close()


@router.get("/import/{task_id}/status")
def import_status(request: Request, task_id: str, _=Depends(verify_token)):
    """轮询导入进度"""
    config = request.app.state.config
    conn = get_connection(config.db_path)
    try:
        task = get_import_task(conn, task_id)
        if not task:
            raise HTTPException(status_code=404, detail={
                "error": {"code": "not_found", "message": f"任务 {task_id} 不存在"}
            })
        return task
    finally:
        conn.close()