"""
知识库管理模块 - 使用 BAAI/bge-small-zh-v1.5 中文嵌入模型（本地离线版）
"""
import os
from pathlib import Path
from typing import List

from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_core.documents import Document

from config import (
    KNOWLEDGE_DIR, CHROMA_DB_DIR, 
    CHUNK_SIZE, CHUNK_OVERLAP, 
    RETRIEVAL_TOP_K
)


class KnowledgeBase:
    """知识库管理类 - 基于本地 BGE 中文嵌入模型"""
    
    def __init__(self):
        self.db_dir = str(CHROMA_DB_DIR)
        self.collection_name = "customer_service_kb"
        
        # 从本地加载 BGE 模型
        print("🔧 正在加载本地 BGE 中文嵌入模型...")
        
        # 直接使用 sentence_transformers 加载本地模型
        from sentence_transformers import SentenceTransformer
        
        # 本地模型路径
        local_model_path = '/home/ubuntu/.cache/huggingface/models--BAAI--bge-small-zh-v1.5/snapshots/'
        # 找到具体的 snapshot 目录
        if os.path.exists(local_model_path):
            snapshots = os.listdir(local_model_path)
            if snapshots:
                model_path = os.path.join(local_model_path, snapshots[0])
            else:
                model_path = 'BAAI/bge-small-zh-v1.5'  # fallback
        else:
            model_path = 'BAAI/bge-small-zh-v1.5'
        
        print(f"   模型路径: {model_path}")
        
        # 加载模型
        self.model = SentenceTransformer(model_path, device='cpu')
        
        # 创建自定义 Embeddings 类
        from langchain_core.embeddings import Embeddings
        
        class LocalBGEEmbeddings(Embeddings):
            def __init__(self, model):
                self.model = model
            
            def embed_documents(self, texts: List[str]) -> List[List[float]]:
                """嵌入文档"""
                embeddings = self.model.encode(texts, normalize_embeddings=True)
                return embeddings.tolist()
            
            def embed_query(self, text: str) -> List[float]:
                """嵌入查询"""
                embedding = self.model.encode(text, normalize_embeddings=True)
                return embedding.tolist()
        
        self.embeddings = LocalBGEEmbeddings(self.model)
        print("✅ BGE 模型加载完成")
        
        # 初始化或加载向量数据库
        self.vectorstore = self._load_or_create_db()
        
    def _load_or_create_db(self) -> Chroma:
        """加载或创建向量数据库"""
        if os.path.exists(self.db_dir) and os.listdir(self.db_dir):
            print(f"📚 加载已有知识库: {self.db_dir}")
            return Chroma(
                collection_name=self.collection_name,
                embedding_function=self.embeddings,
                persist_directory=self.db_dir
            )
        else:
            print(f"🆕 创建新知识库: {self.db_dir}")
            os.makedirs(self.db_dir, exist_ok=True)
            return Chroma(
                collection_name=self.collection_name,
                embedding_function=self.embeddings,
                persist_directory=self.db_dir
            )
    
    def load_documents(self, directory: str = None) -> List[Document]:
        """从目录加载所有文档"""
        directory = directory or str(KNOWLEDGE_DIR)
        documents = []
        
        if not os.path.exists(directory):
            print(f"⚠️ 目录不存在: {directory}")
            return documents
        
        for file_path in Path(directory).glob("*"):
            if file_path.suffix.lower() == ".txt":
                loader = TextLoader(str(file_path), encoding="utf-8")
                documents.extend(loader.load())
                print(f"✅ 加载 TXT: {file_path.name}")
                
            elif file_path.suffix.lower() == ".md":
                loader = TextLoader(str(file_path), encoding="utf-8")
                documents.extend(loader.load())
                print(f"✅ 加载 MD: {file_path.name}")
        
        return documents
    
    def add_documents(self, documents: List[Document]):
        """添加文档到知识库"""
        if not documents:
            print("⚠️ 没有文档需要添加")
            return
        
        # 文本分块
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP,
            separators=["\n\n", "\n", "。", "，", " ", ""]
        )
        chunks = text_splitter.split_documents(documents)
        print(f"📄 文档分块: {len(documents)} 个文档 → {len(chunks)} 个片段")
        
        # 添加到向量数据库
        self.vectorstore.add_documents(chunks)
        print(f"✅ 已添加到知识库")
    
    def search(self, query: str, k: int = None) -> List[Document]:
        """语义搜索相关知识"""
        k = k or RETRIEVAL_TOP_K
        results = self.vectorstore.similarity_search(query, k=k)
        return results
    
    def get_retriever(self):
        """获取检索器（用于 LangChain）"""
        return self.vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_TOP_K})


def init_knowledge_base():
    """初始化知识库（首次运行时使用）"""
    print("🔧 初始化知识库...")
    kb = KnowledgeBase()
    
    # 加载知识库目录中的文档
    documents = kb.load_documents()
    
    if documents:
        kb.add_documents(documents)
        print(f"✅ 知识库初始化完成，共 {len(documents)} 个文档")
    else:
        print("⚠️ 没有找到文档，请在 knowledge_base/ 目录添加文档后重新运行")
    
    return kb


# 全局知识库实例（单例模式）
_knowledge_base = None

def get_knowledge_base() -> KnowledgeBase:
    """获取知识库单例"""
    global _knowledge_base
    if _knowledge_base is None:
        _knowledge_base = KnowledgeBase()
    return _knowledge_base