"""
增强的 Prompt 模板

提供针对不同场景和文件类型优化的 Prompt 模板。
参考 server/aaa/jenius_attachment_knowledge_base/jenius_rag_util.py 和
server/aaa/jenius_personal_knowledge_base/personal_kb_prompt.py 的实现。
"""
from typing import List, Dict, Set


# ==================== 基础 RAG Prompt ====================

RAG_CONTENT_PROMPT = """
## 上传文件内容的分析说明
 - 如果用户问题相关的文件内容可以直接回答用户的问题，则直接基于文件内容回答用户的提问；回答问题不要添加编造成分，不要使用使用大模型的自身知识回答。
 - 输出格式要求：
   * 回答开头应为"根据您上传的文件`related_file_name`", `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果上下文没有提及，`related_file_name`设为空字符串。
   * 如果聊天历史中，系统给出的文件内容并非是用户想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。
 - 如果用户问题相关的文件内容不能回答用户的全部问题，则将文件内容作为上下文，进入工具调用的流程。
 - {rag_content}
## 用户的问题
 - {query}
"""


RAG_FILE_IMAGE_CONTENT_PROMPT = """
## 指令: 根据文件和图片的内容回答用户的问题。
1. 文本文件的问答(docx,pdf,xlsx,xls)：
 - 如果用户问题相关的文件内容可以直接回答用户的问题，则直接基于文件内容回答用户的问题；回答问题不要添加编造成分，不要使用使用大模型的自身知识回答。
 - 如果用户问题相关的文件内容不能回答用户的全部问题，则将文件内容作为上下文，进入工具调用的流程。
 
2. 图片内容的问答(png,jpeg,jpg,bmp)：
 - 如果图片的文字内容为空，或者无法回答用户的问题，进入工具调用流程，使用合适的工具回答用户问题。
 - 如果用户问题为询问图片的主要内容和描述等，进入工具调用流程，使用合适的工具回答用户问题。
 - 如果用户问题涉及图片操作、图片处理、图片加工、进入工具调用流程，使用合适的工具回答用户问题。
 - 如果用户的问题涉及图片的视觉信息，必须进入工具调用流程，使用合适的工具回答用户问题。
   * 视觉信息包括但不限于：物体、场景、颜色、布局、风格、人物、动物、植物、动作。

## 输出格式要求：
- 回答开头应为"根据您上传的文件`related_file_name`", `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果上下文没有提及，`related_file_name`设为空字符串。
- 如果聊天历史中，系统给出的文件内容并非是用户想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。
- 注意：图片问答中，禁止输出"文本内容无法完整回答您的问题"或任何等价说明。

## 输入信息
- 文件/图片内容：{rag_content}
- 用户的问题: {query}
"""


RAG_EXCEL_CONTENT_PROMPT = """
## 上传文件内容的分析说明
 - 如果用户问题相关的文件内容、pandas代码、pandas执行结果,可以直接回答用户的问题，则直接回答用户的提问；回答问题不要添加编造成分，不要使用使用大模型的自身知识回答。
 - 如果用户问题相关的文件内容、pandas代码、pandas执行结果不能回答用户的全部问题，则将他们作为上下文，进入工具调用的流程。
 - 输出格式要求：
   * 回答开头应为"根据您上传的文件`related_file_name`", `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果上下文没有提及，`related_file_name`设为空字符串。
   * 如果聊天历史中，系统给出的文件内容并非是用户想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。
 - {rag_content}
## 用户的问题
 - {query}
"""


# ==================== 知识库 Prompt ====================

KB_CHAT_RAG_PROMPT = """
## 指令：根据文件内容回答用户的问题。请严格按照以下规则处理用户问题：

1. 如果文件内容可以直接回答用户的问题：
- 基于文件内容回答用户的问题。
- 不要添加编造成分，不要使用大模型的自身知识回答。
- 如果文件来源是用户上传的文件，回答开头应为："根据您上传的文件`related_file_name`"
- 如果文件来源是知识库中的文件，回答开头应为："根据您知识库中的文件`related_file_name`"
- 其中`related_file_name`替换为问题相关的文件名,多个文件用英文逗号,分隔,并用markdown反引号`包裹,如果无法确定具体文件名，则不要添加回答开头的这句话。
- 如果聊天历史中，系统给出的文件内容并非是用户想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。

2. 如果文件内容不能回答用户的问题：
 - 不要输出前缀"根据您上传的文件"
 - 将文件内容作为上下文，进入工具调用的流程。

## 知识库的文件内容
 - {kb_rag_content}

## 用户的问题
 - {query}
"""


# ==================== 文件类型特定 Prompt ====================

TEXT_FILE_INSTRUCTION = """
文本文件的问答(docx,pdf,txt)：
- 如果用户问题相关的文件内容可以直接回答用户的问题，则直接基于文件内容回答用户的问题；回答问题不要添加编造成分，不要使用大模型的自身知识回答。
- 如果用户问题相关的文件内容不能回答用户的全部问题，则将文件内容作为上下文，进入工具调用的流程。
"""


IMAGE_FILE_INSTRUCTION = """
图片文件的问答(png,jpeg,jpg,bmp)：
- 如果图片的文字内容为空，或者无法回答用户的问题，进入工具调用流程，使用合适的工具回答用户问题。
- 如果用户问题为询问图片的主要内容和描述等，调用图像理解工具获取更详细的内容来回答用户问题。
- 如果用户问题涉及图片操作、图片处理、图片加工，进入工具调用流程，使用合适的工具回答用户问题。
- 如果用户的问题涉及图片的视觉信息，必须进入工具调用流程，使用合适的工具回答用户问题。
  * 视觉信息包括但不限于：物体、场景、颜色、布局、风格、人物、动物、植物、动作。
- 如果用户要显示图片，则使用file_url进行展示。
- 如果ocr的结果为空或显然无意义，则在回答中不要提及ocr的结果。
"""


EXCEL_FILE_INSTRUCTION = """
表格文件的问答(xlsx,xls,csv)：
- 如果用户问题相关的表格内容可以直接回答用户的问题，则直接基于文件内容回答用户的问题；回答问题不要添加编造成分，不要使用大模型的自身知识回答。
- 如果用户问题相关的表格内容不能回答用户的全部问题，则将文件内容作为上下文，进入工具调用的流程。
- 如果用户问题涉及修改、创建excel表格等操作，进入工具调用流程，使用合适的工具回答用户问题。
"""


AUDIO_FILE_INSTRUCTION = """
音频文件的问答(wav,mp3,flac,m4a,ogg,aac,pcm)：
- 如果用户问题涉及音频文件，进入工具调用流程，使用合适的音频工具回答用户问题。
"""


# ==================== 输出格式 Prompt ====================

CHAT_OUTPUT_FORMAT = """
## 输出格式要求：
- 如果回答的根据来源是用户上传的文件，回答开头应为："根据您上传的文件`related_file_name`"
- 如果需要综合所有文件内容回答，则回答开头根据用户的问题灵活调整
- `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果无法确定具体文件名，则不要添加回答开头的这句话。
- 如果聊天历史中，系统给出的文件内容用户明确表示不是想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。
- 注意：图片问答中，禁止输出"文本内容无法完整回答您的问题"或任何等价说明。

## 输入信息
## 用户上传的文件内容
 - {rag_content}
## 用户的问题
 - {query}
"""


KB_OUTPUT_FORMAT = """
## 输出格式要求：
- 如果回答的根据来源是知识库中的文件，且`related_file_name`是文件名，回答开头应为："根据您知识库中的文件`related_file_name`"
- 如果回答的根据来源是知识库中的网页，且`related_file_name`是网页URL，回答开头应为："根据您知识库中的网页`related_file_name`"
- 如果需要综合所有文件内容回答，则回答开头根据用户的问题灵活调整
- `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果无法确定具体文件名，则不要添加回答开头的这句话。
- 如果聊天历史中，系统给出的文件内容用户明确表示不是想要的，回答开头请加入委婉的回应，例如:"抱歉，我刚才可能理解错了，现在改为分析您需要的文件`related_file_name`。
- 注意：图片问答中，禁止输出"文本内容无法完整回答您的问题"或任何等价说明。

## 输入信息
## 知识库的文件内容
 - {kb_rag_content}
## 用户的问题
 - {query}
"""


# ==================== Prompt 生成函数 ====================

def get_file_extensions(file_list: List[Dict]) -> Set[str]:
    """
    获取文件扩展名集合
    
    Args:
        file_list: 文件列表，每个元素包含 file_name 字段
        
    Returns:
        Set[str]: 文件扩展名集合（小写，带点号）
    """
    extensions = set()
    for file_info in file_list:
        file_name = file_info.get('file_name', '')
        if '.' in file_name:
            ext = '.' + file_name.split('.')[-1].lower()
            extensions.add(ext)
    return extensions


def build_rag_prompt(
    query: str,
    rag_content: str,
    file_list: List[Dict],
    intent_type: str = "summary"
) -> str:
    """
    构建 RAG Prompt
    
    Args:
        query: 用户查询
        rag_content: RAG 内容
        file_list: 文件列表
        intent_type: 意图类型 (summary, excel_analysis, search)
        
    Returns:
        str: 完整的 Prompt
    """
    extensions = get_file_extensions(file_list)
    
    # 根据文件类型选择指令
    instructions = []
    
    if extensions & {'.docx', '.pdf', '.txt'}:
        instructions.append(TEXT_FILE_INSTRUCTION)
    
    if extensions & {'.png', '.jpeg', '.jpg', '.bmp'}:
        instructions.append(IMAGE_FILE_INSTRUCTION)
    
    if extensions & {'.xlsx', '.xls', '.csv'}:
        instructions.append(EXCEL_FILE_INSTRUCTION)
    
    if extensions & {'.wav', '.mp3', '.flac', '.m4a', '.ogg', '.aac', '.pcm'}:
        instructions.append(AUDIO_FILE_INSTRUCTION)
    
    # 根据意图类型选择基础 Prompt
    if intent_type == "excel_analysis":
        base_prompt = RAG_EXCEL_CONTENT_PROMPT
    elif extensions & {'.png', '.jpeg', '.jpg', '.bmp'}:
        base_prompt = RAG_FILE_IMAGE_CONTENT_PROMPT
    else:
        base_prompt = RAG_CONTENT_PROMPT
    
    # 组装完整 Prompt
    full_instructions = "\n".join(instructions) if instructions else ""
    
    if full_instructions:
        # 如果有文件类型特定指令，插入到基础 Prompt 之前
        final_prompt = "## 指令: 根据文件内容回答用户的问题。\n\n" + full_instructions + "\n" + CHAT_OUTPUT_FORMAT
    else:
        final_prompt = base_prompt
    
    return final_prompt.format(query=query, rag_content=rag_content)


def build_kb_rag_prompt(
    query: str,
    kb_rag_content: str,
    file_list: List[Dict]
) -> str:
    """
    构建知识库 RAG Prompt
    
    Args:
        query: 用户查询
        kb_rag_content: 知识库 RAG 内容
        file_list: 文件列表
        
    Returns:
        str: 完整的 Prompt
    """
    extensions = get_file_extensions(file_list)
    
    # 根据文件类型选择指令
    instructions = []
    
    if extensions & {'.docx', '.pdf', '.txt'}:
        instructions.append(TEXT_FILE_INSTRUCTION)
    
    if extensions & {'.png', '.jpeg', '.jpg', '.bmp'}:
        instructions.append(IMAGE_FILE_INSTRUCTION)
    
    if extensions & {'.xlsx', '.xls', '.csv'}:
        instructions.append(EXCEL_FILE_INSTRUCTION)
    
    # 组装完整 Prompt
    full_instructions = "\n".join(instructions) if instructions else ""
    
    if full_instructions:
        final_prompt = "## 指令: 根据知识库文件内容回答用户的问题。\n\n" + full_instructions + "\n" + KB_OUTPUT_FORMAT
    else:
        final_prompt = KB_CHAT_RAG_PROMPT
    
    return final_prompt.format(query=query, kb_rag_content=kb_rag_content)


def build_mixed_rag_prompt(
    query: str,
    chat_rag_content: str,
    kb_rag_content: str,
    chat_file_list: List[Dict],
    kb_file_list: List[Dict]
) -> str:
    """
    构建混合 RAG Prompt（同时包含聊天文件和知识库文件）
    
    Args:
        query: 用户查询
        chat_rag_content: 聊天文件 RAG 内容
        kb_rag_content: 知识库 RAG 内容
        chat_file_list: 聊天文件列表
        kb_file_list: 知识库文件列表
        
    Returns:
        str: 完整的 Prompt
    """
    chat_extensions = get_file_extensions(chat_file_list)
    kb_extensions = get_file_extensions(kb_file_list)
    all_extensions = chat_extensions | kb_extensions
    
    # 根据文件类型选择指令
    instructions = []
    
    if all_extensions & {'.docx', '.pdf', '.txt'}:
        instructions.append(TEXT_FILE_INSTRUCTION)
    
    if all_extensions & {'.png', '.jpeg', '.jpg', '.bmp'}:
        instructions.append(IMAGE_FILE_INSTRUCTION)
    
    if all_extensions & {'.xlsx', '.xls', '.csv'}:
        instructions.append(EXCEL_FILE_INSTRUCTION)
    
    if all_extensions & {'.wav', '.mp3', '.flac', '.m4a', '.ogg', '.aac', '.pcm'}:
        instructions.append(AUDIO_FILE_INSTRUCTION)
    
    # 混合输出格式
    MIXED_OUTPUT_FORMAT = """
## 输出格式要求：
- 如果文件来源是用户上传的文件，回答开头应为："根据您上传的文件`related_file_name`"
- 如果文件来源是知识库中的文件，回答开头应为："根据您知识库中的文件`related_file_name`"
- 如果需要综合所有文件内容回答，则回答开头根据用户的问题灵活调整
- `related_file_name`替换为问题相关的文件名，多个文件用英文逗号,分隔，并用markdown反引号`包裹。如果无法确定具体文件名，则不要添加回答开头的这句话。
- 注意：图片问答中，禁止输出"文本内容无法完整回答您的问题"或任何等价说明。

## 输入信息
## 知识库的文件内容
 - {kb_rag_content}
## 用户上传的文件内容
 - {chat_rag_content}
## 用户的问题
 - {query}
"""
    
    # 组装完整 Prompt
    full_instructions = "\n".join(instructions) if instructions else ""
    
    final_prompt = "## 指令: 根据文件内容回答用户的问题。\n\n" + full_instructions + "\n" + MIXED_OUTPUT_FORMAT
    
    return final_prompt.format(
        query=query,
        chat_rag_content=chat_rag_content,
        kb_rag_content=kb_rag_content
    )