# -*- coding: utf-8 -*-
"""
交互式中文 RAG 系统（支持多次问答 + 流式输出）
作者：AI助手
功能：
  - 加载本地文档（.txt / .json）
  - 智能中文分片（句子边界 + 滑动窗口）
  - BGE 嵌入 + FAISS 索引
  - BGE 重排
  - Gemini 流式生成（支持 1.5-flash 回退）
  - 交互式 CLI：连续问答，输入 quit/exit 退出
"""
import re
import json
import sys
from typing import List, Tuple, Generator

import faiss
import numpy as np

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import google.generativeai as genai

# ======================
# 1. 配置 Google API Key
# ======================
GOOGLE_API_KEY = "your key"  # ←←← 请在此处填写你的 API Key
genai.configure(api_key=GOOGLE_API_KEY)

if GOOGLE_API_KEY == "YOUR_GOOGLE_API_KEY":
    print("❌ 请先在代码中设置你的 Google API Key！")
    sys.exit(1)

# ======================
# 2. 文档加载与中文智能分片
# ======================


def load_documents(file_path: str) -> str:
    """加载 .txt 或 .json 文档"""
    with open(file_path, 'r', encoding='utf-8') as f:
        if file_path.endswith('.json'):
            data = json.load(f)
            if isinstance(data, list):
                text = "\n".join([item.get("text", "") for item in data])
            else:
                text = data.get("text", "")
        else:
            text = f.read()
    return text


def chinese_sentence_split(text: str) -> List[str]:
    """使用标点符号分割中文句子"""
    sentence_endings = re.compile(r'[。！？.!?]+')
    sentences = [s.strip() for s in sentence_endings.split(text) if s.strip()]
    return sentences


def smart_chunking(text: str, chunk_size: int = 512, overlap: int = 64) -> List[str]:
    """基于句子的智能分片，避免词语断裂"""
    sentences = chinese_sentence_split(text)
    chunks = []
    current = ""
    for sent in sentences:
        if len(current) + len(sent) <= chunk_size:
            current += sent + "。"
        else:
            if current:
                chunks.append(current)
            current = sent + "。"
            if len(current) > chunk_size:
                chunks.append(current[:chunk_size])
                current = current[chunk_size:]
    if current:
        chunks.append(current)
    
    # 添加滑动窗口重叠
    if overlap > 0 and len(chunks) > 1:
        overlapped = []
        for i in range(len(chunks)):
            if i == 0:
                overlapped.append(chunks[i])
            else:
                prev_tail = chunks[i-1][-overlap:] if len(chunks[i-1]) > overlap else chunks[i-1]
                overlapped.append(prev_tail + chunks[i])
        return overlapped
    return chunks

# ======================
# 3. 向量存储（BGE 中文嵌入）
# ======================


class VectorStore:
    def __init__(self):
        print("📥 正在加载 BGE 嵌入模型 (BAAI/bge-small-zh-v1.5)...")
        self.embedding_model = SentenceTransformer('BAAI/bge-small-zh-v1.5')
        self.index = None
        self.chunks = []

    def add_documents(self, documents: List[str]):
        self.chunks = documents
        print(f"🧠 正在为 {len(documents)} 个文本块生成嵌入向量...")
        embeddings = self.embedding_model.encode(
            documents,
            batch_size=32,
            show_progress_bar=True,
            normalize_embeddings=True
        )
        embeddings = np.array(embeddings).astype('float32')
        faiss.normalize_L2(embeddings)
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)
        self.index.add(embeddings)
        print("✅ 向量索引构建完成！")

    def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
        query = f"为这个句子生成表示以用于检索相关文章: {query}"
        query_vec = self.embedding_model.encode([query], normalize_embeddings=True).astype('float32')
        faiss.normalize_L2(query_vec)
        scores, indices = self.index.search(query_vec, top_k)
        return [(self.chunks[i], float(scores[0][j])) for j, i in enumerate(indices[0])]

# ======================
# 4. 重排器（BGE Reranker）
# ======================


class ReRanker:
    def __init__(self):
        print("📥 正在加载 BGE 重排模型 (BAAI/bge-reranker-base)...")
        self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-base')
        self.model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
        self.model.eval()
        print("✅ 重排模型加载完成！")

    def rerank(self, query: str, candidates: List[str], top_k: int = 3) -> List[str]:
        pairs = [[query, doc] for doc in candidates]
        with torch.no_grad():
            inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
            scores = self.model(**inputs).logits.view(-1).sigmoid()
        idx = torch.argsort(scores, descending=True)[:top_k]
        return [candidates[i] for i in idx]

# ======================
# 5. 流式生成答案（Gemini）
# ======================


def generate_answer_stream(query: str, retrieved_docs: List[str]) -> Generator[str, None, None]:
    context = "\n\n".join(retrieved_docs)
    prompt = f"""你是一个专业的中文助手，请严格根据以下上下文回答问题。如果上下文没有相关信息，请回答“根据提供的资料无法回答”。

上下文：
{context}

问题：
{query}

回答："""

    model_name = 'gemini-1.5-flash'
    try:
        model = genai.GenerativeModel(model_name)
        response_stream = model.generate_content(prompt, stream=True)
        for chunk in response_stream:
            if chunk.text:
                yield chunk.text
        return  # 成功则退出
    except Exception as e:
        print(f"⚠️ 模型 {model_name} 调用失败，尝试下一个: {e}")


# ======================
# 6. RAG 系统主类
# ======================


class RAGSystem:
    def __init__(self, doc_path: str):
        print(f"📂 正在加载文档: {doc_path}")
        raw_text = load_documents(doc_path)
        self.chunks = smart_chunking(raw_text, chunk_size=512, overlap=64)
        print(f"✂️  文档已分片为 {len(self.chunks)} 个块")

        self.vector_store = VectorStore()
        self.vector_store.add_documents(self.chunks)

        self.reranker = ReRanker()

    def get_answer_stream(self, question: str, top_k_retrieve: int = 10, top_k_rerank: int = 3) -> Generator[str, None, None]:
        # 1. 召回
        retrieved = self.vector_store.search(question, top_k=top_k_retrieve)
        retrieved_texts = [text for text, _ in retrieved]
        # 2. 重排
        reranked_docs = self.reranker.rerank(question, retrieved_texts, top_k=top_k_rerank)
        # 3. 流式生成
        yield from generate_answer_stream(question, reranked_docs)

# ======================
# 7. 交互式主循环
# ======================


def main():
    # 创建示例文档（若不存在）
    DOC_PATH = "knowledge_base.txt"
    # 初始化 RAG 系统
    print("\n🚀 正在初始化 RAG 系统...\n")
    rag = RAGSystem(DOC_PATH)

    # 交互循环
    print("\n💬 RAG 问答系统已启动！")
    print("👉 输入你的问题（输入 'quit' 或 'exit' 退出）\n")

    while True:
        try:
            question = input("❓ 你的问题: ").strip()
            if not question:
                continue
            if question.lower() in ['quit', 'exit', '退出']:
                print("👋 再见！")
                break

            print("🤖 回答: ", end="", flush=True)
            for token in rag.get_answer_stream(question):
                print(token, end="", flush=True)
            print("\n")  # 换行

        except KeyboardInterrupt:
            print("\n\n👋 用户中断，再见！")
            break
        except Exception as e:
            print(f"\n❌ 系统错误: {e}\n")


if __name__ == "__main__":
    main()