"""
简易小说 RAG
三个核心阶段：索引、检索、生成
"""
import os
import re
from openai import OpenAI
import chromadb

# ============================================================
# 配置
# ============================================================∫

NOVEL_PATH = "/Users/chanler/Downloads/西游记.txt" # 小说路径
CHUNK_SIZE = 800 # 每个分块的字符数
OVERLAP = 150 # 相邻分块的重叠字符数(向前重叠)
DISTANCE_METRIC = "cosine"  # 距离度量：cosine / l2 / ip
CHROMA_DB_PATH = "./chroma_db"  # ChromaDB 存储路径

# 自动生成 collection 名称：文件名_度量方式
NOVEL_NAME = "xiyouji"
COLLECTION_NAME = f"{NOVEL_NAME}_{DISTANCE_METRIC}"

# OpenAI 客户端
client = OpenAI(
    base_url="http://127.0.0.1:1234/v1",
    api_key="lm-studio"
)
EMBEDDING_MODEL = "text-embedding-qwen3-embedding-0.6b"
LLM_MODEL = "qwen3-8b"


# ============================================================
# 第一阶段：索引 (Indexing)
# ============================================================

def load_and_split(path, max_chapters=100):
    """
    读取小说并切分为 chunks
    切分策略：先按回目划分章节，章节内重叠切分
    """
    with open(path, 'r', encoding='utf-8') as f:
        text = f.read()

    # 匹配回目标题
    chapter_pattern = r'(第[一二三四五六七八九十百零\d]+回\s+[^\n]+)'
    matches = list(re.finditer(chapter_pattern, text))

    chunks = []
    for i, match in enumerate(matches[:max_chapters]):
        chapter_title = match.group(1).strip()
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        content = text[start:end].strip()

        # 滑动窗口切分
        pos = 0
        while pos < len(content):
            chunk_text = content[pos:pos + CHUNK_SIZE]
            if chunk_text.strip():
                chunks.append((chunk_text.strip(), chapter_title))
            pos += CHUNK_SIZE - OVERLAP

    return chunks


def get_embedding(text):
    """获取文本向量"""
    response = client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=text
    )
    return response.data[0].embedding


def build_index(chunks):
    """构建向量索引"""
    db = chromadb.PersistentClient(path=CHROMA_DB_PATH)

    try:
        db.delete_collection(COLLECTION_NAME)
    except:
        pass

    collection = db.create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": DISTANCE_METRIC}
    )

    print(f"索引 {len(chunks)} 个文本块...")
    for i, (text, chapter) in enumerate(chunks):
        if i % 100 == 0:
            print(f"  进度: {i}/{len(chunks)}")

        collection.add(
            ids=[f"chunk_{i}"],
            embeddings=[get_embedding(text)],
            documents=[text],
            metadatas=[{"chapter": chapter}]
        )

    print(f"完成，共 {collection.count()} 条")
    return collection


# ============================================================
# 第二阶段：检索 (Retrieval)
# ============================================================

def retrieve(collection, query, top_k=3):
    """检索最相关的文本块"""
    results = collection.query(
        query_embeddings=[get_embedding(query)],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )
    return results


# ============================================================
# 第三阶段：生成 (Generation)
# ============================================================

def build_prompt(query, contexts):
    """构建增强提示词"""
    context_text = "\n\n".join(
        f"【{meta['chapter']}】\n{doc}"
        for doc, meta in contexts
    )

    return f"""你是一个小说问答助手。请根据以下参考内容回答问题。

要求：
1. 只根据参考内容作答，不要编造
2. 如果没有相关信息，明确说明
3. 标注出处（回目）

参考内容：
{context_text}

问题：{query}

回答："""


def generate(prompt):
    """调用 LLM 生成回答"""
    response = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
        max_tokens=1000
    )
    return response.choices[0].message.content


def rag_query(collection, query, top_k=1):
    """完整 RAG 流程：检索 -> 增强 -> 生成"""
    results = retrieve(collection, query, top_k)

    contexts = list(zip(
        results["documents"][0],
        results["metadatas"][0]
    ))

    prompt = build_prompt(query, contexts)
    print(f"\n{'='*50}\n【完整提示词】\n{'='*50}\n{prompt}\n")
    answer = generate(prompt)

    return {
        "answer": answer,
        "sources": results["metadatas"][0],
        "distances": results["distances"][0]
    }


# ============================================================
# 索引管理
# ============================================================

def get_collection():
    """获取已有索引"""
    db = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    return db.get_collection(COLLECTION_NAME)


def index_exists():
    """检查索引是否存在"""
    db = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    try:
        return db.get_collection(COLLECTION_NAME).count() > 0
    except:
        return False


# ============================================================
# 主流程
# ============================================================

if __name__ == "__main__":
    import sys

    force_index = "--index" in sys.argv

    print("=" * 50)
    print("简易小说 RAG")
    print("=" * 50)
    print(f"Collection: {COLLECTION_NAME}")

    if force_index or not index_exists():
        print("\n[索引阶段]")
        chunks = load_and_split(NOVEL_PATH)
        print(f"切分为 {len(chunks)} 个文本块")
        collection = build_index(chunks)
    else:
        print("\n[加载已有索引]")
        collection = get_collection()
        print(f"已加载 {collection.count()} 条记录")

    print("\n[检索与生成]")
    query = "菩萨告诉孙悟空有人救他的原文句子或段落，并标注回目"
    print(f"查询: {query}\n")

    result = rag_query(collection, query, top_k=2)

    print("="*50)
    print("【参考】")
    for src, dist in zip(result["sources"], result["distances"]):
        print(f" {src['chapter']} (相似度: {1-dist:.1%})")
    print("="*50)
    print("【回答】")
    print(result["answer"])
    print("="*50)