HuiBrain/app/api/v1/import_export.py

"""
导入导出 API 路由
提供文件导入和知识库导出接口
支持 .md / .txt / .docx 格式
"""

from __future__ import annotations

import json
import os
from datetime import datetime
from typing import Optional

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.database import get_db
from app.services.import_service import ImportService

router = APIRouter()

# 支持的导入文件格式
SUPPORTED_EXTENSIONS = (".md", ".txt", ".docx")


@router.post("/file", summary="从文件导入知识")
async def import_from_file(
    file: UploadFile = File(..., description="上传文件（.md / .txt / .docx）"),
    course_name: Optional[str] = Form(None, description="课程名称"),
    teacher_name: Optional[str] = Form(None, description="讲师名称"),
    live_date: Optional[str] = Form(None, description="直播日期 (YYYY-MM-DD)"),
    db: AsyncSession = Depends(get_db),
):
    """
    上传文件并导入到知识库。

    - Markdown (.md)：解析 frontmatter 元数据，按 ## 标题分页
    - 纯文本 (.txt)：整体作为一个页面
    - Word 文档 (.docx)：提取文本，按标题分页
    - 导入后自动分块、生成嵌入向量、建立全文索引
    """
    # 验证文件类型
    if not file.filename:
        raise HTTPException(status_code=400, detail="文件名不能为空")

    suffix = os.path.splitext(file.filename)[1].lower()
    if suffix not in SUPPORTED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"仅支持 {', '.join(SUPPORTED_EXTENSIONS)} 文件",
        )

    # 保存上传文件到 transcripts 目录
    dest_dir = settings.transcripts_dir
    dest_path = dest_dir / file.filename

    # 如果文件已存在，添加序号
    counter = 1
    original_stem = os.path.splitext(file.filename)[0]
    while dest_path.exists():
        dest_path = dest_dir / f"{original_stem}_{counter}{suffix}"
        counter += 1

    # 写入文件
    content = await file.read()
    dest_path.write_bytes(content)

    # 执行导入
    service = ImportService(db)
    try:
        result = await service.import_file(
            file_path=str(dest_path),
            course_name=course_name,
            teacher_name=teacher_name,
            live_date=live_date,
        )
        return {
            "message": "导入成功",
            "file": str(dest_path),
            **result,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")


@router.post("/directory", summary="从目录批量导入")
async def import_from_directory(
    directory: str = Form(..., description="数据目录中的子目录名（相对于 transcripts 目录）"),
    course_name: Optional[str] = Form(None, description="课程名称"),
    teacher_name: Optional[str] = Form(None, description="讲师名称"),
    live_date: Optional[str] = Form(None, description="直播日期"),
    db: AsyncSession = Depends(get_db),
):
    """
    批量导入指定目录下的所有文件（.md / .txt / .docx）。
    """
    target_dir = settings.transcripts_dir / directory
    if not target_dir.exists():
        raise HTTPException(status_code=404, detail=f"目录不存在: {directory}")

    service = ImportService(db)
    total_pages = 0
    total_chunks = 0
    files_processed = 0
    errors: list[str] = []

    for filepath in sorted(target_dir.iterdir()):
        if filepath.suffix.lower() not in SUPPORTED_EXTENSIONS:
            continue

        try:
            result = await service.import_file(
                file_path=str(filepath),
                course_name=course_name,
                teacher_name=teacher_name,
                live_date=live_date,
            )
            total_pages += result["pages"]
            total_chunks += result["chunks"]
            files_processed += 1
        except Exception as e:
            errors.append(f"{filepath.name}: {str(e)}")

    return {
        "message": "批量导入完成",
        "files_processed": files_processed,
        "total_pages": total_pages,
        "total_chunks": total_chunks,
        "errors": errors if errors else None,
    }


@router.get("/export", summary="导出知识库")
async def export_knowledge_base(
    format: str = "json",
    course_name: Optional[str] = None,
    db: AsyncSession = Depends(get_db),
):
    """
    导出知识库内容。

    Args:
        format: 导出格式，支持 json / markdown
        course_name: 按课程名称过滤导出
    """
    from sqlalchemy import text

    query = text("""
        SELECT id, title, content, source_file, course_name, teacher_name,
               live_date, page_number, metadata_json, created_at
        FROM knowledge_pages
        WHERE (:course_name IS NULL OR course_name = :course_name)
        ORDER BY created_at DESC
    """)
    result = await db.execute(query, {"course_name": course_name})
    pages = result.fetchall()

    if format == "markdown":
        # 导出为 Markdown 文件
        lines = []
        for row in pages:
            lines.append(f"# {row.title}")
            lines.append("")
            meta_parts = []
            if row.course_name:
                meta_parts.append(f"课程: {row.course_name}")
            if row.teacher_name:
                meta_parts.append(f"讲师: {row.teacher_name}")
            if row.live_date:
                meta_parts.append(f"日期: {row.live_date}")
            if meta_parts:
                lines.append(" | ".join(meta_parts))
                lines.append("")
            lines.append(row.content)
            lines.append("")
            lines.append("---")
            lines.append("")

        export_content = "\n".join(lines)
        ext = "md"
    else:
        # 导出为 JSON
        pages_data = []
        for row in pages:
            pages_data.append({
                "id": row.id,
                "title": row.title,
                "content": row.content,
                "source_file": row.source_file,
                "course_name": row.course_name,
                "teacher_name": row.teacher_name,
                "live_date": row.live_date,
                "page_number": row.page_number,
                "metadata_json": row.metadata_json,
                "created_at": str(row.created_at),
            })
        export_content = json.dumps(pages_data, ensure_ascii=False, indent=2)
        ext = "json"

    # 保存到 exports 目录
    export_dir = settings.exports_dir
    export_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"export_{timestamp}.{ext}"
    export_path = export_dir / filename
    export_path.write_text(export_content, encoding="utf-8")

    return {
        "message": "导出成功",
        "file": filename,
        "pages": len(pages),
        "format": format,
    }