# ingest.py
import os
from langchain_community.document_loaders import (
    DirectoryLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
    PyPDFLoader
)
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 配置路径
DOCS_DIR = "./docs"
PERSIST_DIR = "./chroma_db"

def load_documents():
    """加载所有支持的文档"""
    documents = []

    # 加载 .txt 文件
    txt_loader = DirectoryLoader(DOCS_DIR, glob="*.txt", loader_cls=TextLoader, loader_kwargs={"encoding": "utf-8"})
    documents.extend(txt_loader.load())

    # 加载 .md 文件
    md_loader = DirectoryLoader(DOCS_DIR, glob="*.md", loader_cls=UnstructuredMarkdownLoader)
    documents.extend(md_loader.load())

    # 加载 .pdf 文件（可选）
    # pdf_loader = DirectoryLoader(DOCS_DIR, glob="*.pdf", loader_cls=PyPDFLoader)
    # documents.extend(pdf_loader.load())

    print(f"✅ 共加载 {len(documents)} 个文档片段")
    return documents

def split_documents(documents):
    """将文档切分为小块"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        length_function=len,
        is_separator_regex=False,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"✅ 切分为 {len(chunks)} 个文本块")
    return chunks

if __name__ == "__main__":
    docs = load_documents()
    chunks = split_documents(docs)

    # 保存到临时文件（用于下一步嵌入）
    os.makedirs(PERSIST_DIR, exist_ok=True)
    with open(os.path.join(PERSIST_DIR, "chunks.txt"), "w", encoding="utf-8") as f:
        for i, chunk in enumerate(chunks):
            f.write(f"--- Chunk {i} ---\n")
            f.write(chunk.page_content + "\n\n")
    print("✅ 文本块已保存（供调试）")

# build_vector_db.py
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from ingest import load_documents, split_documents
import os

# 嵌入模型配置（完全本地）
embedding_model_name = "BAAI/bge-small-zh-v1.5"
model_kwargs = {'device': 'cpu'}          # 若有 GPU 可改为 'cuda'
encode_kwargs = {'normalize_embeddings': True}  # BGE 推荐设置

embeddings = HuggingFaceBgeEmbeddings(
    model_name=embedding_model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："  # BGE 中文模型推荐 query 指令
)

PERSIST_DIR = "./chroma_db"

def build_vector_db():
    print("🔍 正在加载文档...")
    docs = load_documents()
    chunks = split_documents(docs)

    print("🧠 正在生成向量并构建数据库（首次运行会下载模型到本地缓存）...")
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=PERSIST_DIR,
        collection_name="my_knowledge"
    )
    print(f"✅ 向量数据库已保存到 {PERSIST_DIR}")

if __name__ == "__main__":
    build_vector_db()

# query.py
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddings

embedding_model_name = "BAAI/bge-small-zh-v1.5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}

embeddings = HuggingFaceBgeEmbeddings(
    model_name=embedding_model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："
)

vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
    collection_name="my_knowledge"
)

def search(query: str, k: int = 3):
    docs = vectorstore.similarity_search(query, k=k)
    print(f"\n🔍 问题：{query}\n")
    for i, doc in enumerate(docs):
        print(f"【结果 {i+1}】（来源: {doc.metadata.get('source', '未知')}）")
        print(doc.page_content[:500] + "...\n")
    return docs

if __name__ == "__main__":
    while True:
        question = input("\n请输入你的问题（输入 'quit' 退出）：")
        if question.lower() == 'quit':
            break
        search(question)

# rag_query.py
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# === 配置 ===
embedding_model_name = "BAAI/bge-small-zh-v1.5"
llm_model = "qwen:4b"  # Ollama 中的模型名

# === 嵌入与向量库 ===
embeddings = HuggingFaceBgeEmbeddings(
    model_name=embedding_model_name,
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True},
    query_instruction="为这个句子生成表示以用于检索相关文章："
)

vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
    collection_name="my_knowledge"
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# === LLM 与 Prompt ===
llm = OllamaLLM(model=llm_model, temperature=0.3)

template = """
你是一个知识助手，请根据以下提供的上下文信息回答问题。
只使用上下文中的内容，不要编造。如果上下文没有相关信息，请回答“根据现有资料无法回答”。

上下文：
{context}

问题：{question}
答案：
"""

prompt = ChatPromptTemplate.from_template(template)

# === RAG 链 ===
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# === 查询 ===
if __name__ == "__main__":
    while True:
        question = input("\n🧠 请输入你的问题（'quit' 退出）：")
        if question.lower() == 'quit':
            break
        print("\n🤔 正在思考...\n")
        answer = rag_chain.invoke(question)
        print(f"✅ 答案：{answer}\n")

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('BAAI/bge-small-zh-v1.5')
model.save('./bge-small-zh')  # 保存到本地

# auto_updater.py
import os
import time
import hashlib
from pathlib import Path
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from build_vector_db import build_vector_db  # 复用之前写的构建函数

DOCS_DIR = "./docs"
STATE_FILE = "./chroma_db/update_state.txt"

def compute_docs_hash():
    """计算 docs/ 目录下所有文件的 hash，用于判断是否变化"""
    hash_md5 = hashlib.md5()
    paths = sorted(Path(DOCS_DIR).rglob("*"))
    for path in paths:
        if path.is_file():
            stat = path.stat()
            hash_md5.update(str(stat.st_mtime).encode())
            hash_md5.update(str(stat.st_size).encode())
    return hash_md5.hexdigest()

def load_last_hash():
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, "r") as f:
            return f.read().strip()
    return None

def save_current_hash():
    current_hash = compute_docs_hash()
    os.makedirs(os.path.dirname(STATE_FILE), exist_ok=True)
    with open(STATE_FILE, "w") as f:
        f.write(current_hash)

class DocUpdateHandler(FileSystemEventHandler):
    def __init__(self, callback):
        self.callback = callback
        self.last_trigger = 0

    def on_any_event(self, event):
        # 防抖：5 秒内只触发一次
        if time.time() - self.last_trigger > 5:
            if event.src_path.endswith(('.txt', '.md', '.pdf')):
                print(f"\n📁 检测到文档变更: {event.src_path}")
                self.last_trigger = time.time()
                self.callback()

def start_watcher():
    """启动文件监控"""
    observer = Observer()
    event_handler = DocUpdateHandler(on_docs_changed)
    observer.schedule(event_handler, DOCS_DIR, recursive=True)
    observer.start()
    print(f"👀 正在监控 {os.path.abspath(DOCS_DIR)} 目录...")

    # 初始检查是否需要首次构建
    if not os.path.exists("./chroma_db/chroma.sqlite3"):  # ChromaDB 默认文件
        print("🆕 首次运行：正在构建向量库...")
        build_vector_db()
        save_current_hash()
    else:
        current_hash = compute_docs_hash()
        last_hash = load_last_hash()
        if current_hash != last_hash:
            print("🔄 检测到历史变更：正在重建向量库...")
            build_vector_db()
            save_current_hash()

    return observer

def on_docs_changed():
    """文档变更时的回调"""
    build_vector_db()
    save_current_hash()
    print("✅ 向量库已更新！")

# build_vector_db.py（修改版）
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from ingest import load_documents, split_documents
import os
import shutil

PERSIST_DIR = "./chroma_db"
embedding_model_name = "BAAI/bge-small-zh-v1.5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
query_instruction = "为这个句子生成表示以用于检索相关文章："

# 全局嵌入模型（避免重复加载）
_EMBEDDING_MODEL = None

def get_embedding_model():
    global _EMBEDDING_MODEL
    if _EMBEDDING_MODEL is None:
        _EMBEDDING_MODEL = HuggingFaceBgeEmbeddings(
            model_name=embedding_model_name,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs,
            query_instruction=query_instruction
        )
    return _EMBEDDING_MODEL

def build_vector_db():
    """构建向量数据库（覆盖写入）"""
    # 删除旧数据库（Chroma 不支持高效更新）
    if os.path.exists(PERSIST_DIR):
        shutil.rmtree(PERSIST_DIR)
        os.makedirs(PERSIST_DIR)

    docs = load_documents()
    chunks = split_documents(docs)

    embeddings = get_embedding_model()
    print("🧠 正在生成向量...")
    Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=PERSIST_DIR,
        collection_name="my_knowledge"
    )
    print("✅ 向量库重建完成")

# app.py
import streamlit as st
from langchain_chroma import Chroma
from build_vector_db import get_embedding_model
from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
import os

# === 配置 ===
PERSIST_DIR = "./chroma_db"
USE_RAG = True  # 设为 False 则只返回检索结果

# === 初始化 ===
@st.cache_resource
def get_retriever():
    embeddings = get_embedding_model()
    vectorstore = Chroma(
        persist_directory=PERSIST_DIR,
        embedding_function=embeddings,
        collection_name="my_knowledge"
    )
    return vectorstore.as_retriever(search_kwargs={"k": 3})

@st.cache_resource
def get_rag_chain():
    retriever = get_retriever()
    llm = OllamaLLM(model="qwen:4b", temperature=0.3)
    template = """
你是一个知识助手，请根据以下提供的上下文信息回答问题。
只使用上下文中的内容，不要编造。如果上下文没有相关信息，请回答“根据现有资料无法回答”。

上下文：
{context}

问题：{question}
答案：
"""
    prompt = ChatPromptTemplate.from_template(template)
    return (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

# === Streamlit UI ===
st.set_page_config(page_title="我的本地知识库", layout="wide")
st.title("🧠 我的本地知识库（完全离线）")

# 输入框
query = st.text_input("请输入你的问题：", placeholder="例如：项目的关键时间节点是什么？")

if query:
    with st.spinner("正在检索和思考..."):
        if USE_RAG:
            try:
                answer = get_rag_chain().invoke(query)
                st.subheader("✅ 答案")
                st.write(answer)
            except Exception as e:
                st.error(f"模型调用失败（请确认 Ollama 正在运行）：{str(e)}")
                USE_RAG = False  # 回退到仅检索
        
        # 显示检索到的原文（始终显示）
        retriever = get_retriever()
        docs = retriever.invoke(query)
        st.subheader("📚 相关原文片段")
        for i, doc in enumerate(docs):
            source = doc.metadata.get("source", "未知")
            st.markdown(f"**来源 {i+1}:** `{os.path.basename(source)}`")
            st.text_area("", doc.page_content, height=150, key=f"doc_{i}")

# 侧边栏
with st.sidebar:
    st.header("⚙️ 系统状态")
    db_exists = os.path.exists(os.path.join(PERSIST_DIR, "chroma.sqlite3"))
    st.success("🟢 向量库已加载") if db_exists else st.warning("🔴 向量库未构建")
    
    ollama_status = "🟢 Ollama 可用（若启用 RAG）" if USE_RAG else "⚪ 仅检索模式"
    st.write(ollama_status)
    
    st.info("📁 文档目录: `./docs`")
    st.info("🔒 所有数据仅保存在本地")

# ingest.py （更新版：支持 PDF / DOCX / TXT / MD，完全本地）

import os
from langchain_community.document_loaders import (
    DirectoryLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
)
from langchain_community.document_loaders.unstructured import UnstructuredFileLoader
from typing import List
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="unstructured")

DOCS_DIR = "./docs"
PERSIST_DIR = "./chroma_db"

def load_documents() -> List:
    """加载所有支持的文档（TXT, MD, PDF, DOCX），完全本地解析"""
    documents = []

    # === .txt 文件 ===
    txt_loader = DirectoryLoader(
        DOCS_DIR,
        glob="*.txt",
        loader_cls=TextLoader,
        loader_kwargs={"encoding": "utf-8"}
    )
    documents.extend(txt_loader.load())

    # === .md 文件 ===
    md_loader = DirectoryLoader(
        DOCS_DIR,
        glob="*.md",
        loader_cls=UnstructuredMarkdownLoader,
        loader_kwargs={"encoding": "utf-8"}
    )
    documents.extend(md_loader.load())

    # === .pdf 文件 ===
    pdf_files = [str(f) for f in Path(DOCS_DIR).glob("*.pdf")]
    for pdf_file in pdf_files:
        print(f"📄 正在解析 PDF: {os.path.basename(pdf_file)}")
        loader = UnstructuredFileLoader(
            pdf_file,
            mode="single",  # 或 "elements"（更细粒度）
            strategy="fast",  # "hi_res" 更准但慢，需 OCR
            # 关键：强制本地，不调用 API
            api_url=None,   # 禁用远程 API
            # 启用本地 OCR（如果 PDF 是扫描件）
            ocr_languages=["chi_sim", "eng"],  # 中英混合
        )
        documents.extend(loader.load())

    # === .docx 文件 ===
    docx_files = [str(f) for f in Path(DOCS_DIR).glob("*.docx")]
    for docx_file in docx_files:
        print(f"📝 正在解析 DOCX: {os.path.basename(docx_file)}")
        loader = UnstructuredFileLoader(
            docx_file,
            mode="single",
            strategy="fast",
            api_url=None,  # 禁用远程 API
        )
        documents.extend(loader.load())

    print(f"✅ 共加载 {len(documents)} 个文档片段")
    return documents

# --- 保留原有的 split_documents 函数 ---
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pathlib import Path

def split_documents(documents):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        length_function=len,
        is_separator_regex=False,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"✅ 切分为 {len(chunks)} 个文本块")
    return chunks

if __name__ == "__main__":
    docs = load_documents()
    chunks = split_documents(docs)

    os.makedirs(PERSIST_DIR, exist_ok=True)
    with open(os.path.join(PERSIST_DIR, "chunks.txt"), "w", encoding="utf-8") as f:
        for i, chunk in enumerate(chunks):
            f.write(f"--- Chunk {i} ---\n")
            f.write(chunk.page_content + "\n\n")
    print("✅ 文本块已保存（供调试）")

# 在 UnstructuredFileLoader 中显式指定 OCR 选项
loader = UnstructuredFileLoader(
    pdf_file,
    strategy="hi_res",          # 启用高级布局分析
    ocr_languages="chi_sim+eng", # 中英混合
    include_page_breaks=False,
    api_url=None,
)