"""
RAG完整实现：从文档到问答
- 文档加载与分割
- 向量嵌入
- 向量存储与检索
- LLM生成回答
"""
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from typing import List, Tuple
import os

# 1. 文档加载
class DocumentLoader:
    """文档加载器"""
    
    @staticmethod
    def load_directory(path: str, glob_pattern: str = "**/*.txt") -> List:
        """加载目录下所有文档"""
        loader = DirectoryLoader(path, glob=glob_pattern, loader_cls=TextLoader)
        documents = loader.load()
        return documents
    
    @staticmethod
    def load_file(file_path: str) -> List:
        """加载单个文档"""
        loader = TextLoader(file_path, encoding="utf-8")
        return loader.load()

# 2. 文档分割
class TextSplitter:
    """文本分割器"""
    
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""]
        )
    
    def split(self, documents) -> List:
        """分割文档"""
        return self.splitter.split_documents(documents)

# 3. 向量存储
class VectorStore:
    """向量存储器"""
    
    def __init__(self, embedding_model=None):
        self.embedding = embedding_model or OpenAIEmbeddings()
        self.vectorstore = None
    
    def create(self, documents: List, save_path: str = None):
        """创建向量存储"""
        self.vectorstore = FAISS.from_documents(documents, self.embedding)
        
        if save_path:
            self.vectorstore.save_local(save_path)
            print(f"✅ 向量库已保存到: {save_path}")
    
    def load(self, path: str):
        """加载向量存储"""
        self.vectorstore = FAISS.load_local(path, self.embedding)
    
    def search(self, query: str, k: int = 4) -> List:
        """相似度检索"""
        if not self.vectorstore:
            raise ValueError("向量库未初始化")
        return self.vectorstore.similarity_search(query, k=k)
    
    def search_with_score(self, query: str, k: int = 4) -> List[Tuple]:
        """带分数的检索"""
        return self.vectorstore.similarity_search_with_score(query, k=k)

# 4. RAG链
class RAGChain:
    """RAG问答链"""
    
    def __init__(self, llm, vectorstore: VectorStore):
        self.llm = llm
        self.vectorstore = vectorstore
        
        # 提示词模板
        self.prompt_template = """基于以下上下文回答问题。如果上下文中没有相关信息，请说明你不知道。

上下文:
{context}

问题: {question}

回答:"""
        
        self.prompt = ChatPromptTemplate.from_template(self.prompt_template)
        self.chain = None
        self.build_chain()
    
    def build_chain(self):
        """构建RAG链"""
        retriever = self.vectorstore.vectorstore.as_retriever()
        
        self.chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | self.prompt
            | self.llm
            | StrOutputParser()
        )
    
    def invoke(self, question: str) -> str:
        """执行问答"""
        return self.chain.invoke(question)

# 5. 混合检索增强
class HybridRetriever:
    """混合检索器 - 向量+关键词"""
    
    def __init__(self, vectorstore: FAISS):
        self.vectorstore = vectorstore
        self.search_kwargs = {"k": 10}
    
    def search(self, query: str) -> List:
        """混合检索"""
        # 1. 向量检索
        vector_results = self.vectorstore.similarity_search(query, k=5)
        
        # 2. MMR检索（增加多样性）
        mmr_results = self.vectorstore.max_marginal_relevance_search(
            query, k=3, fetch_k=10
        )
        
        # 3. 合并去重
        seen = set()
        combined = []
        for doc in vector_results + mmr_results:
            doc_id = doc.metadata.get("source", "") + doc.page_content[:50]
            if doc_id not in seen:
                seen.add(doc_id)
                combined.append(doc)
        
        return combined[:4]


# 使用示例
def main():
    # 准备测试文档
    os.makedirs("docs", exist_ok=True)
    with open("docs/ai_intro.txt", "w", encoding="utf-8") as f:
        f.write("""人工智能(AI)是一门研究如何让机器具有智能的学科。
机器学习是AI的核心技术之一，让计算机从数据中学习规律。
深度学习是机器学习的一个分支，使用神经网络模型。
大语言模型(LLM)是深度学习的应用，可以理解和生成人类语言。
RAG技术可以增强LLM的知识，让AI回答私有知识库的问题。
""")
    
    with open("docs/rag_guide.txt", "w", encoding="utf-8") as f:
        f.write("""RAG(检索增强生成)是一种结合检索和生成的技术。
RAG的工作流程：加载文档 -> 分割文本 -> 向量化 -> 存储 -> 检索 -> 生成。
RAG的优点：1) 可以使用最新信息 2) 减少幻觉 3) 可以引用来源。
常用的向量数据库有FAISS、Chroma、Milvus、Pinecone等。
""")
    
    # 1. 加载文档
    print("📄 加载文档...")
    documents = DocumentLoader.load_directory("docs")
    print(f"  加载了 {len(documents)} 个文档")
    
    # 2. 分割文档
    print("✂️ 分割文档...")
    splitter = TextSplitter(chunk_size=100, chunk_overlap=20)
    chunks = splitter.split(documents)
    print(f"  分割成 {len(chunks)} 个块")
    
    # 3. 创建向量存储
    print("🔢 创建向量索引...")
    embedding = OpenAIEmbeddings()
    vectorstore = VectorStore(embedding)
    vectorstore.create(chunks, save_path="vectorstore")
    
    # 4. 测试检索
    print("\n🔍 测试检索...")
    results = vectorstore.search("什么是RAG？", k=2)
    for i, doc in enumerate(results):
        print(f"  结果{i+1}: {doc.page_content[:100]}...")
    
    # 5. 构建RAG链
    print("\n🤖 构建RAG问答链...")
    llm = ChatOpenAI(model="gpt-4", temperature=0)
    vectorstore.load("vectorstore")
    rag = RAGChain(llm, vectorstore)
    
    # 6. 执行问答
    questions = [
        "什么是RAG？",
        "人工智能和机器学习有什么关系？",
        "常用的向量数据库有哪些？"
    ]
    
    print("\n💬 RAG问答测试:")
    for q in questions:
        print(f"\n问题: {q}")
        answer = rag.invoke(q)
        print(f"回答: {answer}")

main()