course-agent-od/backend/src/services/knowledge_base_service.py

"""
知识库管理服务
"""
import logging
import os
import hashlib
from pathlib import Path
from typing import List, Dict, Any, Optional
from datetime import datetime
from sqlalchemy.orm import Session
from sqlalchemy import and_

from ..models.document import Document, DocumentChunk
from ..models.knowledge_base import KnowledgeBase
from ..models.user import User
from ..core.config import get_settings
from .document_service import DocumentService

logger = logging.getLogger(__name__)

settings = get_settings()


class KnowledgeBaseService:
    """知识库管理服务"""
    
    def __init__(self, db: Session):
        self.db = db
        self.document_service = DocumentService(db)
        self.knowledge_base_dir = Path(settings.knowledge_base_dir)
        
    def scan_directory(self, directory: Optional[str] = None) -> Dict[str, Any]:
        """扫描知识库目录并批量导入"""
        if directory is None:
            directory = self.knowledge_base_dir
        else:
            directory = Path(directory)

        if not directory.exists():
            return {"success": False, "message": f"目录不存在: {directory}"}

        # 先确保每个子目录都有对应的系统知识库
        self.ensure_system_knowledge_bases()

        results = {
            "scanned_files": 0,
            "new_files": 0,
            "updated_files": 0,
            "skipped_files": 0,
            "errors": []
        }

        # 递归扫描目录
        for file_path in directory.rglob("*"):
            if file_path.is_file() and self._is_supported_file(file_path):
                try:
                    result = self.process_file(str(file_path))
                    results["scanned_files"] += 1

                    if result["status"] == "new":
                        results["new_files"] += 1
                    elif result["status"] == "updated":
                        results["updated_files"] += 1
                    elif result["status"] == "skipped":
                        results["skipped_files"] += 1
                    elif result["status"] == "error":
                        results["errors"].append({
                            "file": str(file_path),
                            "error": result["error"]
                        })

                except Exception as e:
                    results["errors"].append({
                        "file": str(file_path),
                        "error": str(e)
                    })

        results["success"] = len(results["errors"]) == 0
        return results

    def ensure_system_knowledge_bases(self):
        """确保 knowledge_base_dir 下每个子目录都有对应的系统知识库"""
        if not self.knowledge_base_dir.exists():
            return

        for subdir in self.knowledge_base_dir.iterdir():
            if subdir.is_dir():
                self._get_or_create_system_kb(subdir.name)

    def _get_or_create_system_kb(self, name: str) -> int:
        """按名称查找或创建系统知识库，返回 knowledge_base_id"""
        kb = self.db.query(KnowledgeBase).filter(
            KnowledgeBase.name == name,
            KnowledgeBase.is_system == True
        ).first()

        if kb:
            return kb.id

        # 检查是否已有同名非系统KB，如有则升级为系统KB
        existing = self.db.query(KnowledgeBase).filter(
            KnowledgeBase.name == name,
            KnowledgeBase.is_system == False
        ).first()
        if existing:
            existing.is_system = True
            admin = self.db.query(User).filter(User.is_superuser == True).first()
            if admin:
                existing.user_id = admin.id
            self.db.commit()
            self.db.refresh(existing)
            logger.info(f"升级为系统知识库: {name} (id={existing.id})")
            return existing.id

        # 找到 admin 用户（或任意 superuser）作为 owner
        admin = self.db.query(User).filter(User.is_superuser == True).first()
        if not admin:
            admin = self.db.query(User).first()

        kb = KnowledgeBase(
            name=name,
            description=f"系统知识库：{name}",
            user_id=admin.id if admin else 1,
            is_system=True
        )
        self.db.add(kb)
        self.db.commit()
        self.db.refresh(kb)
        logger.info(f"自动创建系统知识库: {name} (id={kb.id})")
        return kb.id

    def _resolve_knowledge_base(self, file_path: Path) -> Optional[int]:
        """从文件路径解析对应的系统知识库 ID

        data/knowledge_base/国土空间规划文献/paper.pdf → 知识库 "国土空间规划文献"
        """
        try:
            relative = file_path.relative_to(self.knowledge_base_dir)
            parts = relative.parts
            if len(parts) >= 2:
                subdir_name = parts[0]
                return self._get_or_create_system_kb(subdir_name)
        except ValueError:
            pass
        return None

    def process_file(self, file_path: str) -> Dict[str, Any]:
        """处理单个文件（检查、提取、入库）"""
        try:
            file_path = Path(file_path)
            
            # 检查文件是否存在
            if not file_path.exists():
                return {"status": "error", "error": "文件不存在"}
            
            # 检查文件类型
            if not self._is_supported_file(file_path):
                return {"status": "skipped", "message": "不支持的文件类型"}
            
            # 获取文件信息
            file_stat = file_path.stat()
            file_size = file_stat.st_size
            last_modified = datetime.fromtimestamp(file_stat.st_mtime)
            
            # 计算文件哈希
            file_hash = self._calculate_file_hash(file_path)
            
            # 检查文件是否已存在
            existing_doc = self.db.query(Document).filter(
                and_(
                    Document.file_path == str(file_path),
                    Document.source_type == "knowledge_base"
                )
            ).first()
            
            if existing_doc:
                # 检查是否需要更新
                if (existing_doc.last_modified and
                    existing_doc.last_modified >= last_modified and
                    existing_doc.file_hash == file_hash):
                    return {"status": "skipped", "message": "文件未修改"}

                # 如果文档没有关联知识库，尝试关联
                if existing_doc.knowledge_base_id is None:
                    kb_id = self._resolve_knowledge_base(file_path)
                    if kb_id:
                        existing_doc.knowledge_base_id = kb_id
                        self.db.commit()

                # 更新现有文档
                return self._update_document(existing_doc, file_path, file_size, last_modified, file_hash)
            else:
                # 解析知识库 ID
                kb_id = self._resolve_knowledge_base(file_path)
                # 创建新文档
                return self._create_document(file_path, file_size, last_modified, file_hash, knowledge_base_id=kb_id)
                
        except Exception as e:
            return {"status": "error", "error": str(e)}
    
    def _is_supported_file(self, file_path: Path) -> bool:
        """检查是否为支持的文件类型"""
        return file_path.suffix.lower() in settings.allowed_extensions
    
    def _calculate_file_hash(self, file_path: Path) -> str:
        """计算文件哈希"""
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()
    
    def _create_document(self, file_path: Path, file_size: int, last_modified: datetime, file_hash: str, 
                        knowledge_base_id: Optional[int] = None, user_id: Optional[int] = None) -> Dict[str, Any]:
        """创建新文档"""
        try:
            # 计算相对路径用于描述
            try:
                relative_path = file_path.relative_to(self.knowledge_base_dir)
            except ValueError:
                relative_path = file_path.name
            
            # 创建文档记录
            document = Document(
                user_id=user_id,  # 如果指定了user_id则使用，否则为None（系统文档）
                knowledge_base_id=knowledge_base_id,  # 关联知识库
                filename=file_path.name,
                original_filename=file_path.name,
                file_path=str(file_path),
                file_size=file_size,
                file_type=file_path.suffix.lower(),
                title=file_path.stem,
                description=f"知识库文档: {relative_path}",
                is_processed=False,
                is_public=True,
                source_type="knowledge_base",
                last_modified=last_modified,
                file_hash=file_hash
            )
            
            self.db.add(document)
            self.db.commit()
            self.db.refresh(document)
            
            # 处理文档（向量化）- 在后台异步处理，不阻塞主流程
            try:
                import threading
                import asyncio as _asyncio
                def process_in_background():
                    try:
                        loop = _asyncio.new_event_loop()
                        loop.run_until_complete(self.document_service.process_document(document.id))
                        loop.close()
                    except Exception as e:
                        logger.error(f"后台处理文档 {document.id} 失败: {e}")

                thread = threading.Thread(target=process_in_background, daemon=True)
                thread.start()
            except Exception as e:
                logger.error(f"启动文档处理线程失败: {e}")

            return {"status": "new", "document_id": document.id, "message": "文档创建成功，正在处理中"}
                
        except Exception as e:
            self.db.rollback()
            return {"status": "error", "error": str(e)}
    
    
    def _update_document(self, document: Document, file_path: Path, file_size: int, last_modified: datetime, file_hash: str) -> Dict[str, Any]:
        """更新现有文档"""
        try:
            # 更新文档信息
            document.file_size = file_size
            document.last_modified = last_modified
            document.file_hash = file_hash
            document.is_processed = False  # 标记为未处理，需要重新处理
            
            # 删除旧的文档块
            self.db.query(DocumentChunk).filter(DocumentChunk.document_id == document.id).delete()
            
            self.db.commit()
            
            # 重新处理文档 - 在后台异步处理
            try:
                import threading
                import asyncio as _asyncio
                def process_in_background():
                    try:
                        loop = _asyncio.new_event_loop()
                        loop.run_until_complete(self.document_service.process_document(document.id))
                        loop.close()
                    except Exception as e:
                        logger.error(f"后台处理文档 {document.id} 失败: {e}")

                thread = threading.Thread(target=process_in_background, daemon=True)
                thread.start()
            except Exception as e:
                logger.error(f"启动文档处理线程失败: {e}")

            return {"status": "updated", "document_id": document.id, "message": "文档更新成功，正在重新处理中"}
                
        except Exception as e:
            self.db.rollback()
            return {"status": "error", "error": str(e)}
    
    def is_file_updated(self, file_path: str) -> bool:
        """检查文件是否需要更新"""
        try:
            file_path = Path(file_path)
            if not file_path.exists():
                return False
            
            # 获取文件信息
            file_stat = file_path.stat()
            last_modified = datetime.fromtimestamp(file_stat.st_mtime)
            file_hash = self._calculate_file_hash(file_path)
            
            # 查询数据库中的记录
            existing_doc = self.db.query(Document).filter(
                and_(
                    Document.file_path == str(file_path),
                    Document.source_type == "knowledge_base"
                )
            ).first()
            
            if not existing_doc:
                return True  # 新文件
            
            # 检查修改时间和哈希
            if (existing_doc.last_modified and 
                existing_doc.last_modified < last_modified):
                return True
            
            if existing_doc.file_hash != file_hash:
                return True
            
            return False
            
        except Exception:
            return True  # 出错时默认需要更新
    
    def get_knowledge_base_status(self) -> Dict[str, Any]:
        """获取知识库状态"""
        try:
            # 统计文档数量
            total_docs = self.db.query(Document).filter(
                Document.source_type == "knowledge_base"
            ).count()
            
            processed_docs = self.db.query(Document).filter(
                and_(
                    Document.source_type == "knowledge_base",
                    Document.is_processed == True
                )
            ).count()
            
            # 统计文件大小
            total_size = self.db.query(Document).filter(
                Document.source_type == "knowledge_base"
            ).with_entities(Document.file_size).all()
            
            total_size_bytes = sum(size[0] for size in total_size) if total_size else 0
            
            # 统计文件类型
            file_types = {}
            docs_by_type = self.db.query(Document.file_type).filter(
                Document.source_type == "knowledge_base"
            ).all()
            
            for file_type in docs_by_type:
                file_type_str = file_type[0]
                file_types[file_type_str] = file_types.get(file_type_str, 0) + 1
            
            return {
                "total_documents": total_docs,
                "processed_documents": processed_docs,
                "unprocessed_documents": total_docs - processed_docs,
                "total_size_bytes": total_size_bytes,
                "total_size_mb": round(total_size_bytes / (1024 * 1024), 2),
                "file_types": file_types,
                "knowledge_base_dir": str(self.knowledge_base_dir),
                "directory_exists": self.knowledge_base_dir.exists()
            }
            
        except Exception as e:
            return {"error": str(e)}
    
    def delete_document(self, document_id: int) -> Dict[str, Any]:
        """删除知识库文档"""
        try:
            document = self.db.query(Document).filter(
                and_(
                    Document.id == document_id,
                    Document.source_type == "knowledge_base"
                )
            ).first()
            
            if not document:
                return {"success": False, "message": "文档不存在"}
            
            # 删除文档块
            self.db.query(DocumentChunk).filter(DocumentChunk.document_id == document_id).delete()
            
            # 删除文档记录
            self.db.delete(document)
            self.db.commit()
            
            return {"success": True, "message": "文档删除成功"}
            
        except Exception as e:
            self.db.rollback()
            return {"success": False, "message": str(e)}
    
    def reindex_document(self, document_id: int) -> Dict[str, Any]:
        """重新索引指定文档"""
        try:
            document = self.db.query(Document).filter(
                and_(
                    Document.id == document_id,
                    Document.source_type == "knowledge_base"
                )
            ).first()
            
            if not document:
                return {"success": False, "message": "文档不存在"}
            
            # 删除旧的文档块
            self.db.query(DocumentChunk).filter(DocumentChunk.document_id == document_id).delete()
            
            # 重新处理文档
            import asyncio as _asyncio
            success = _asyncio.get_event_loop().run_until_complete(
                self.document_service.process_document(document_id)
            )
            
            if success:
                return {"success": True, "message": "文档重新索引成功"}
            else:
                return {"success": False, "message": "文档重新索引失败"}
                
        except Exception as e:
            self.db.rollback()
            return {"success": False, "message": str(e)}
    
    def import_files_to_knowledge_base(
        self, 
        knowledge_base_id: int, 
        directory: Path, 
        user_id: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        从指定目录导入真实文件到知识库
        
        Args:
            knowledge_base_id: 知识库ID
            directory: 要扫描的目录路径
            user_id: 用户ID（可选，用于系统知识库时可为None）
        
        Returns:
            包含统计信息的字典
        """
        if isinstance(directory, str):
            directory = Path(directory)
        
        if not directory.exists():
            return {
                "success": False,
                "message": f"目录不存在: {directory}",
                "scanned_files": 0,
                "new_files": 0,
                "updated_files": 0,
                "skipped_files": 0,
                "errors": []
            }
        
        results = {
            "scanned_files": 0,
            "new_files": 0,
            "updated_files": 0,
            "skipped_files": 0,
            "errors": []
        }
        
        # 递归扫描目录下的所有真实文件
        all_files = [f for f in directory.rglob("*") if f.is_file() and self._is_supported_file(f)]
        total_files = len(all_files)
        
        logger.info(f"找到 {total_files} 个支持的文件，开始处理...")
        
        for idx, file_path in enumerate(all_files, 1):
            try:
                # 获取文件信息
                file_stat = file_path.stat()
                file_size = file_stat.st_size
                last_modified = datetime.fromtimestamp(file_stat.st_mtime)
                file_hash = self._calculate_file_hash(file_path)
                
                # 检查文件是否已存在于该知识库中
                existing_doc = self.db.query(Document).filter(
                    and_(
                        Document.file_path == str(file_path),
                        Document.knowledge_base_id == knowledge_base_id,
                        Document.source_type == "knowledge_base"
                    )
                ).first()
                
                results["scanned_files"] += 1
                
                # 显示进度（每10个文件或最后一个文件时显示）
                if idx % 10 == 0 or idx == total_files:
                    percentage = (idx * 100) // total_files if total_files > 0 else 0
                    logger.info(f"处理进度: {idx}/{total_files} ({percentage}%)")
                
                if existing_doc:
                    # 检查是否需要更新
                    if (existing_doc.last_modified and 
                        existing_doc.last_modified >= last_modified and
                        existing_doc.file_hash == file_hash):
                        results["skipped_files"] += 1
                        continue
                    
                    # 更新现有文档
                    result = self._update_document(existing_doc, file_path, file_size, last_modified, file_hash)
                    if result["status"] == "updated":
                        results["updated_files"] += 1
                    else:
                        results["errors"].append({
                            "file": str(file_path),
                            "error": result.get("error", "更新失败")
                        })
                else:
                    # 创建新文档
                    result = self._create_document(
                        file_path, 
                        file_size, 
                        last_modified, 
                        file_hash,
                        knowledge_base_id=knowledge_base_id,
                        user_id=user_id
                    )
                    if result["status"] == "new":
                        results["new_files"] += 1
                    else:
                        results["errors"].append({
                            "file": str(file_path),
                            "error": result.get("error", "创建失败")
                        })
                        
            except Exception as e:
                results["errors"].append({
                    "file": str(file_path),
                    "error": str(e)
                })
        
        logger.info("文件处理完成")
        results["success"] = len(results["errors"]) == 0
        return results