from milvus_client import connect_to_milvus, create_collection, insert_documents, search_documents
from document_processor import DocumentProcessor
from config import LLM_MODEL
from org.config import SALT_PROMPT, REISASOL_PROMPT


class RAGSystem:
    def __init__(self):
        # 连接Milvus
        connect_to_milvus()

        # 创建集合
        self.collection = create_collection()

        # 初始化文档处理器
        self.processor = DocumentProcessor()

        # 加载集合
        self.collection.load()

    def add_documents(self, documents):
        """添加文档到知识库"""
        # 分割文档
        split_texts = self.processor.split_documents(documents)

        # 生成嵌入
        embeddings = self.processor.generate_embeddings(split_texts)

        # 准备插入数据
        texts = [text['content'] for text in split_texts]
        sources = [text['source'] for text in split_texts]

        # 插入数据
        insert_documents(self.collection, embeddings, texts, sources)

        return len(split_texts)

    def query(self, question, top_k=3):
        """基础查询：生成标准答案"""
        # 生成查询嵌入
        query_embedding = self.processor.embedder.embed_query(question)

        # 检索相关文档
        retrieved_docs = search_documents(self.collection, query_embedding, top_k)

        # 构造上下文
        context = "\n".join([doc['text'] for doc in retrieved_docs])

        # 生成标准答案
        answer = self._generate_answer(question, context)

        return {
            'answer': answer,
            'retrieved_docs': retrieved_docs,
            'context': context
        }

    def convert_songs_to_documents(songs_data):
        """将歌曲数据转换为文档格式"""
        documents = []
        for song in songs_data:
            content = f"歌曲名称: {song['title']}, 歌手: {song['artist']}, BPM: {song['bpm']}, 版本: {song['version']}"
            documents.append({
                'content': content,
                'source': f"歌曲数据 - {song['title']}"
            })
        return documents

    def add_song_data(self, songs_data):
        """添加歌曲数据到知识库"""
        # 转换数据格式
        song_documents = self.convert_songs_to_documents(songs_data)

        # 复用现有文档处理流程
        split_texts = self.processor.split_documents(song_documents)
        embeddings = self.processor.generate_embeddings(split_texts)

        texts = [text['content'] for text in split_texts]
        sources = [text['source'] for text in split_texts]

        insert_documents(self.collection, embeddings, texts, sources)
        return len(split_texts)

    # -------------------------- 新增：角色扮演式查询方法 --------------------------
    def role_play_query(self, question, role, top_k=3):
        """
        角色扮演查询：按指定角色生成符合身份的答案
        :param question: 用户问题（如“Python怎么定义函数？”）
        :param role: 指定角色（如“Python讲师”“产品经理”“小学生辅导员”）
        :param top_k: 检索相关文档的数量
        :return: 包含角色化答案的结果字典
        """
        # 1. 复用现有检索逻辑（和query方法完全一致，确保上下文相关性）
        query_embedding = self.processor.embedder.embed_query(question)
        retrieved_docs = search_documents(self.collection, query_embedding, top_k)
        context = "\n".join([doc['text'] for doc in retrieved_docs])

        # 2. 调用新增的“角色化答案生成方法”（区别于基础的_generate_answer）
        role_answer = self._generate_role_answer(question, context, role)

        # 3. 返回结构和query方法一致，便于后续使用
        return {
            'answer': role_answer,  # 角色化的答案
            'role': role,           # 明确返回当前角色
            'retrieved_docs': retrieved_docs,
            'context': context
        }
    # ----------------------------------------------------------------------------

    def _generate_answer(self, question, context):
        """基础答案生成：无角色限制"""
        from config import OLLAMA_BASE_URL
        import requests

        prompt = f"""
        基于以下上下文回答问题。如果上下文不包含相关信息，请说明无法基于提供的资料回答。
        要求：答案简洁、准确，符合技术文档规范。

        上下文:
        {context}

        问题: {question}

        回答:
        """

        try:
            base_url = OLLAMA_BASE_URL.rstrip('/')
            ollama_api_url = f"{base_url}/api/generate"
            request_body = {
                "model": LLM_MODEL,
                "prompt": prompt,
                "stream": False
            }

            response = requests.post(ollama_api_url, json=request_body, timeout=30)
            if response.status_code == 200:
                return response.json().get("response", "未获取到答案内容")
            else:
                return f"Ollama API请求失败，状态码：{response.status_code}，原因：{response.text}"

        except requests.exceptions.ConnectionError:
            return f"连接Ollama服务失败（{ollama_api_url}），请检查网络"
        except requests.exceptions.Timeout:
            return f"连接Ollama服务超时（30秒），可能是模型生成过慢"
        except Exception as e:
            return f"生成答案出错：{str(e)}"

    # -------------------------- 新增：角色化答案生成方法 --------------------------
    def _generate_role_answer(self, question, context, role):
        """
        角色化答案生成：按指定角色调整prompt语气和内容风格
        :param question: 用户问题
        :param context: 检索到的相关文档上下文
        :param role: 指定角色
        :return: 符合角色身份的答案
        """
        from config import OLLAMA_BASE_URL
        import requests

        # 核心区别：在prompt中加入“角色定义”，让模型按身份生成内容
        # 不同角色对应不同的语气要求（可根据需要扩展更多角色的提示词）
        role_prompt_map = {
            "纱露朵": SALT_PROMPT,
            "Reisasol" : REISASOL_PROMPT
       }

        # 获取当前角色的提示词（若角色不在预设中，用默认提示词）
        role_desc = role_prompt_map.get(
            role,
            f"你是{role}，请基于上下文回答问题，保持语气符合{role}的身份。"
        )

        # 构造角色化prompt
        prompt = f"""
        {role_desc}
        核心要求：基于以下上下文回答,如果上下文不包含相关信息，请给出最合适的回答，并确保内容符合{role}的身份。”。

        上下文:
        {context}

        问题: {question}

        回答:
        """

        # 后续API调用逻辑和_generate_answer一致（复用网络请求代码）
        try:
            base_url = OLLAMA_BASE_URL.rstrip('/')
            ollama_api_url = f"{base_url}/api/generate"
            request_body = {
                "model": LLM_MODEL,
                "prompt": prompt,
                "stream": False
            }

            response = requests.post(ollama_api_url, json=request_body, timeout=30)
            if response.status_code == 200:
                return response.json().get("response", "未获取到答案内容")
            else:
                return f"Ollama API请求失败，状态码：{response.status_code}，原因：{response.text}"

        except requests.exceptions.ConnectionError:
            return f"连接Ollama服务失败（{ollama_api_url}），请检查网络"
        except requests.exceptions.Timeout:
            return f"连接Ollama服务超时（30秒），可能是模型生成过慢"
        except Exception as e:
            return f"生成角色化答案出错：{str(e)}"
    # ----------------------------------------------------------------------------