import requests
import json
import dashscope
from http import HTTPStatus
from typing import List, Dict

# ================= 配置区域 =================
# 1. 设置 DashScope API Key (这里填入你提供的 Key)
dashscope.api_key = "sk-8b091493de594c5e9eb42f12f1cc5805"

# 2. 本地后端地址 (刚才写的 FastAPI)
BACKEND_SEARCH_URL = "http://127.0.0.1:8000/api/v2/search"

# 3. 选择模型 (qwen-turbo, qwen-plus, qwen-max)
MODEL_NAME = dashscope.Generation.Models.qwen_plus
# ===========================================

class WikiBot:
    def __init__(self):
        self.history = [] #以此保存多轮对话上下文(可选)

    def search_knowledge_base(self, query: str, top_k: int = 5) -> List[Dict]:
        """调用本地后端接口检索相关知识"""
        try:
            payload = {
                "query": query,
                "limit": top_k
            }
            # 调用 /api/v2/search，后端会自动做 embedding
            resp = requests.post(BACKEND_SEARCH_URL, json=payload)
            
            if resp.status_code == 200:
                data = resp.json()
                if data.get("code") == 1:
                    return data.get("data", [])
            
            print(f"[Warning] 检索失败: {resp.text}")
            return []
        except Exception as e:
            print(f"[Error] 连接后端失败: {e}")
            return []

    def build_prompt(self, query: str, context_chunks: List[Dict]) -> str:
        """构建 RAG 提示词"""
        
        if not context_chunks:
            return f"用户问题：{query}\n\n当前知识库中没有找到相关信息，请直接告知用户无法回答。"

        # 拼接参考资料
        context_str = ""
        for idx, item in enumerate(context_chunks):
            # 这里把 source_url 也带上，方便 AI 引用来源
            source = item.get('source_url', '未知来源')
            content = item.get('content', '').strip()
            context_str += f"【参考资料 {idx+1}】(来源: {source}):\n{content}\n\n"

        # 系统提示词 (System Prompt)
        prompt = f"""你是一个专业的 Wiki 知识库助手。
请严格根据下方的【参考上下文】来回答用户的【问题】。

要求：
1. 回答要准确、简洁，并整合不同参考资料中的信息。
2. 如果【参考上下文】中包含答案，请用自己的话回答，并在句尾标注来源，例如 [参考资料 1]。
3. 如果【参考上下文】与问题无关或不包含答案，请直接回答：“知识库中暂未收录相关信息”，不要编造答案。
4. 保持回答格式清晰（可以使用 Markdown）。

====== 参考上下文 开始 ======
{context_str}
====== 参考上下文 结束 ======

用户问题：{query}
"""
        return prompt

    def chat(self, query: str):
        """主对话逻辑"""
        print(f"\n🔍 正在检索知识库...")
        
        # 1. 检索
        chunks = self.search_knowledge_base(query)
        print(f"✅ 找到 {len(chunks)} 条相关资料")

        # 2. 构建 Prompt
        prompt = self.build_prompt(query, chunks)
        
        # (可选) 调试时打印 prompt 看看给 AI 喂了什么
        # print(f"DEBUG PROMPT:\n{prompt}\n")

        print("🤖 Wiki助手正在思考...\n" + "-"*30)

        # 3. 调用 DashScope 生成 (流式输出)
        responses = dashscope.Generation.call(
            model=MODEL_NAME,
            messages=[
                {'role': 'system', 'content': 'You are a helpful assistant.'},
                {'role': 'user', 'content': prompt}
            ],
            result_format='message',  # 设置输出为 message 格式
            stream=True,              # 开启流式输出
            incremental_output=True   # 增量输出
        )

        full_content = ""
        for response in responses:
            if response.status_code == HTTPStatus.OK:
                text = response.output.choices[0]['message']['content']
                full_content += text
                print(text, end='', flush=True)
            else:
                print(f"\nRequest id: {response.request_id}, Status code: {response.status_code}, error code: {response.code}, error message: {response.message}")

        print("\n" + "-"*30 + "\n")

# ================= 运行入口 =================
if __name__ == "__main__":
    bot = WikiBot()
    print("✨ Wiki 知识库助手已启动 (输入 'q' 或 'exit' 退出)")
    print("⚠️  请确保后端服务 (main.py) 正在 localhost:8000 运行")
    
    while True:
        user_input = input("\n🙋 请输入问题: ").strip()
        
        if user_input.lower() in ['q', 'exit', 'quit']:
            print("再见！")
            break
            
        if not user_input:
            continue
            
        bot.chat(user_input)