"""
工作流：普通问答Agent（使用LangGraph）
输入：文本（转化后和原文本query）、图片描述、语音文本
步骤：
1. 使用RAG和GraphRAG检索相关信息
2. 生成回答
使用 OpenaiAPI 方法调用大模型，尽可能使用异步
"""
import sys
import os

# 添加项目根目录到Python搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from typing import TypedDict, Optional, Dict, Any, List
from langgraph.graph import StateGraph, START, END
from workflows.history_manager import init_history

from tools.function_tool import (
    graph_rag_search
)
from tools.rag_tools import rag_search
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls, get_current_stream_handler, get_all_callbacks
from utils.image_utils import extract_image_paths, find_closest_image_path, normalize_markdown_images, get_image_prompt_guidance
from prompts import BAIKE_PROMPTS, COMMON_PROMPTS
import asyncio
import json
import re
import random


# =============== 定义状态 ===============
class QAState(TypedDict):
    """普通问答工作流状态"""
    extracted_text: str  # 已提取的文本（统一接口，只接收文本）
    combined_query: str
    retrieval_query: str  # 用于RAG/图谱检索的重写后query
    route_flag: str  # 路由标识
    history_message: str  # 历史消息（JSON格式）
    history_messages: List[Dict[str, Any]]  # 过滤后的历史消息列表
    need_retrieval: Optional[bool]  # 是否需要检索

    # RAG检索结果
    rag_search_result: Optional[List[Dict[str, Any]]]  # RAG检索结果
    graph_search_result: Optional[str]  # 图谱检索结果

    # 最终响应
    response: str  # 最终响应
    actions: List[Dict[str, Any]]  # 按钮信息列表 [{"type": "...", "label": "...", "icon": "..."}]
    suggestedReplies: List[Dict[str, Any]]  # 建议回复问题列表
    error_message: str  # 错误信息


# =============== 节点函数 ===============

def get_baike_history_context(state: QAState) -> Dict[str, Any]:
    """
    知识问答步骤：获取问答类历史信息
    问答工作流只关注历史问答记录，用于提供上下文
    """
    history_message = state.get("history_message", "") or ""
    return init_history(history_message)


async def judge_need_retrieval(state: QAState) -> Dict[str, Any]:
    """
    判断用户问题是否需要检索知识库
    """
    original_query = (state.get("combined_query") or state.get("extracted_text") or "").strip()
    history_messages = state.get("history_messages", []) or []
    
    prompt = BAIKE_PROMPTS["judge_need_retrieval"].format(
        history_messages=history_messages if history_messages else '无',
        original_query=original_query
    )
    
    need_retrieval = True
    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(prompt, model=None)
        if result:
            result = result.strip().lower()
            if result == "false":
                need_retrieval = False
    except Exception as e:
        print(f"判断是否需要检索失败: {e}")
        need_retrieval = True



    print(f"是否需要检索: {need_retrieval}")
    return {"need_retrieval": need_retrieval}


async def rewrite_query_with_history(state: QAState) -> Dict[str, Any]:
    """
    使用历史上下文对当前问题进行query重写，生成适合RAG/图谱检索的独立查询语句。
    只让大模型在有限的最近历史中挑选与当前问题语义相关的内容进行融合，避免简单拼接全量历史。
    """
    import json

    # 原始问题：优先使用 combined_query，其次使用 extracted_text
    original_query = (state.get("combined_query") or state.get("extracted_text") or "").strip()

    history_messages = state.get("history_messages", []) or []
    # 控制历史长度：只保留最近的若干条，避免上下文过长
    trimmed_history = history_messages[-8:] if history_messages else []

    try:
        history_str = json.dumps(trimmed_history, ensure_ascii=False)
    except Exception:
        history_str = str(trimmed_history)

    # 如果没有历史或当前问题过短，就直接跳过重写
    if not original_query:
        return {"retrieval_query": original_query}
    # 在构造 prompt 前，提取上一轮用户 query（如果有）
    last_user_query = ""
    if trimmed_history:
        # 从后往前找最近一条 role == "user" 的消息
        for msg in reversed(trimmed_history):
            if msg.get("role") == "user":
                last_user_query = msg.get("content", "").strip()
                break
    prompt = BAIKE_PROMPTS["rewrite_query_with_history"].format(
        last_user_query=last_user_query or '（无）',
        history_str=history_str,
        original_query=original_query
    )

    rewritten_query = original_query
    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(prompt, model=None)
        if result:
            candidate = result.strip().strip("“”\"'")
            # 只在模型返回的内容明显为单行短文本时替换，避免异常输出
            if candidate and len(candidate) <= 200 and "\n" not in candidate:
                rewritten_query = candidate
    except Exception as e:
        print(f"query 重写失败: {e}")


    title_options = [f"✨正在分析中，请稍等..."]
    start_event = {
        "type": "function_execution",
        "title": random.choice(title_options),
        "details": f"正在调取百科数据..."
    }
    for callback in get_all_callbacks():
        try:
            callback(start_event)
        except Exception:
            pass

    print("原始检索query:", original_query)
    print("重写后检索query:", last_user_query+rewritten_query)

    return {"retrieval_query":last_user_query+rewritten_query}


@track_function_calls
async def search_rag_and_graph(state: QAState) -> Dict[str, Any]:
    """
    进行文本知识检索和图谱知识检索
    """
    kb_id = None
    if state.get("route_flag") == "rules":
        kb_id = '3'

    # 优先使用重写后的检索query，其次使用 combined_query，最后回退到原始 extracted_text
    query = state.get("retrieval_query") or state.get("combined_query") or state.get("extracted_text", "")
    if "用户问题：" in query:
        query = query.split("用户问题：", 1)[1]

    print("Graph 检索query", query)
    extracted_text = state.get("extracted_text", "")
    if "用户问题：" in extracted_text:
        extracted_text = extracted_text.split("用户问题：", 1)[1]

    try:
        # 异步并行调用RAG检索和图谱检索
        print(kb_id)
        print(33333333333333333333333333333333333333)

        async def _do_rag_search():
            if kb_id:
                return await rag_search(
                    query=extracted_text,
                    kb_id=kb_id,
                    top_k=2
                )
            else:
                return await rag_search(
                    query=extracted_text,
                    top_k=2
                )

        async def _do_graph_search():
            return await graph_rag_search.ainvoke({
                "query": query,
                "top_k": 3
            })

        rag_result, graph_results = await asyncio.gather(
            _do_rag_search(),
            _do_graph_search()
        )

        # 处理 RAG 结果，提取成列表格式
        rag_search_result = []
        if rag_result.get("success", False):
            source_citation = rag_result.get("sourceCitation", {})
            for resource, items in source_citation.items():
                for item in items:
                    rag_search_result.append({
                        "text": item.get("text", ""),
                        "filename": item.get("filename", ""),
                        "resource": item.get("resource", ""),
                        "score": item.get("score", 0.0)
                    })

        # 只保留前 4 条数据
        rag_search_result = rag_search_result[:2]

        return {
            "rag_search_result": rag_search_result,
            "graph_search_result": graph_results if isinstance(graph_results, str) else ""
        }
    except Exception as e:
        return {
            "rag_search_result": [],
            "graph_search_result": "",
            "error_message": f"检索失败: {str(e)}"
        }



async def generate_suggested_replies_baike(answer: str, question: str) -> List[Dict[str, Any]]:
    """
    生成针对问答结果的建议回复问题
    使用大模型生成两个相关问题
    """
    prompt = BAIKE_PROMPTS["generate_suggested_replies"].format(
        question=question,
        answer=answer[:500]
    )
    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(prompt, model=None, json_output=True)
        # 尝试解析JSON
        # 提取JSON部分
        json_match = re.search(r'\[.*\]', result, re.DOTALL)
        if json_match:
            json_str = json_match.group(0)
            suggested_replies = json.loads(json_str)
            # 验证格式
            if isinstance(suggested_replies, list) and len(suggested_replies) >= 2:
                # 确保每个元素都有title和content
                formatted_replies = []
                for reply in suggested_replies[:2]:
                    if isinstance(reply, dict) and "title" in reply and "content" in reply:
                        formatted_replies.append({
                            "title": str(reply["title"]),
                            "content": str(reply["content"])
                        })
                if len(formatted_replies) == 2:
                    return formatted_replies

        # 如果解析失败，返回默认问题
        return [
            {"title": "这个设备的具体工作原理是什么？", "content": "这个设备的具体工作原理是什么？"},
            {"title": "还有哪些相关的知识点？", "content": "还有哪些相关的知识点？"}
        ]
    except Exception as e:
        # 如果生成失败，返回默认问题
        return [
            {"title": "这个设备的具体工作原理是什么？", "content": "这个设备的具体工作原理是什么？"},
            {"title": "还有哪些相关的知识点？", "content": "还有哪些相关的知识点？"}
        ]


@track_function_calls
async def generate_answer(state: QAState) -> Dict[str, Any]:
    """
    正在生成问答结果
    """
    extracted_text = state.get("combined_query", "")
    rag_results = state.get("rag_search_result", [])
    graph_results = state.get("graph_search_result", "")
    history_messages = state.get("history_messages", []) or []
    need_retrieval = state.get("need_retrieval", True)
    print("history_message(qa)", history_messages)
    print("是否使用检索结果:", need_retrieval)

    try:
        # 1. 处理检索结果（如果有）
        rag_ctx_parts: List[str] = []
        file_name_list = []
        all_source_text = ""
        original_image_paths = []
        
        if need_retrieval:
            # 收集所有 RAG 结果文本（只取前 4 条）
            for item in (rag_results or [])[:1]:
                print(item)
                print(111111111111111111111111111)
                if isinstance(item, dict):
                    text = (item.get("text") or "").strip()
                    file_name = (item.get("filename") or "").strip()
                    if text:
                        rag_ctx_parts.append(text)
                    if file_name:
                        file_name_list.append(file_name)
            
            # 处理图谱结果
            print("123456", graph_results)
            if graph_results and "图谱检索完成，但未返回有效数据。" in graph_results:
                graph_results = ""
            print("654321", graph_results)
            
            # 从所有结果中提取图片路径
            if graph_results:
                all_source_text = graph_results + "\n" + "\n".join(rag_ctx_parts)
            else:
                all_source_text = "\n".join(rag_ctx_parts)
            original_image_paths = extract_image_paths(all_source_text)
            print("提取到的原始图片路径:", original_image_paths)
            
            # 发送事件通知（可选）
            title_options = [f"✨已综合 {len(rag_ctx_parts)} 条检索结果", f"✨本次检索策略为综合多源信息"]
            start_event = {
                "type": "function_execution",
                "title": random.choice(title_options),
                "details": f"综合了 {len(rag_ctx_parts)} 条检索结果"
            }
            for callback in get_all_callbacks():
                try:
                    callback(start_event)
                except Exception:
                    pass

        query_desc = extracted_text if extracted_text else "当前未明确描述问题"

        # 构建系统提示词
        if state.get("route_flag") == "rules":
            domain_desc = "舰船修理相关的法规、规范、行业标准及技术文件"
        else:
            domain_desc = "工业设备的结构原理、维护规程、故障诊断与维修操作"

        # ========== 第一步：专注于生成高质量内容 ==========
        # 根据是否有检索结果构建不同的提示词
        if need_retrieval and all_source_text.strip():
            content_system_prompt = BAIKE_PROMPTS["generate_answer_with_retrieval_system"].format(domain_desc=domain_desc)
            content_user_content = BAIKE_PROMPTS["generate_answer_with_retrieval_user"].format(
                extracted_text=extracted_text,
                rag_count=len(rag_ctx_parts),
                all_source_text=all_source_text
            )
        else:
            content_system_prompt = BAIKE_PROMPTS["generate_answer_without_retrieval_system"].format(domain_desc=domain_desc)
            content_user_content = BAIKE_PROMPTS["generate_answer_without_retrieval_user"].format(extracted_text=extracted_text)

        content_messages = []
        if history_messages:
            content_messages.extend(history_messages)
        content_messages.append({"role": "user", "content": content_user_content})

        # 第一步：生成内容
        raw_content = await OpenaiAPI.open_api_chat_without_thinking(
            model=None,
            system_prompt=content_system_prompt,
            messages=content_messages,
            temperature=0.5
        )

        # 过滤掉异常标签（如果有）
        raw_content = re.sub(r'ynchroneg>.*?ost switching>', '', raw_content, flags=re.DOTALL)
        raw_content = raw_content.strip() if raw_content else "抱歉，无法生成回答。"

        title_options = ["🤖 诊断分析中", "🤖 故障预检中"]
        start_event = {
            "type": "function_execution",
            "title": random.choice(title_options),
            "details": raw_content[:30] + "..."
        }
        for callback in get_all_callbacks():
            try:
                callback(start_event)
            except Exception:
                pass


        # ========== 第二步：专注于规范格式 ==========
        # 构建第二步的提示词：只关注格式规范，不修改内容
        format_system_prompt = COMMON_PROMPTS["format_system"]
        
        format_user_content = COMMON_PROMPTS["format_user_with_image_examples"].format(raw_content=raw_content)

        # 第二步：规范格式（流式输出）
        stream_handler = get_current_stream_handler()
        answer = ""
        accumulated_content = ""  # 用于逐步发送给前端的内容
        
        # 使用 async for 遍历流式生成器
        async for chunk in OpenaiAPI.open_api_chat_stream(
            model=None,
            system_prompt=format_system_prompt,
            messages=[{"role": "user", "content": format_user_content}],
            temperature=0.1
        ):
            answer += chunk
            accumulated_content += chunk
            # 每收到新的 chunk 就发送一次更新
            if stream_handler:
                stream_handler.send_stream_content(accumulated_content)

        answer = answer.strip() if answer else raw_content

        # LaTeX 公式空格规范化
        def normalize_latex_spacing(text: str) -> str:
            placeholder = "\x00DOUBLE_DOLLAR\x00"
            text = text.replace("$$", placeholder)
            text = re.sub(r'(?<!\s)\$', r' $', text)
            text = re.sub(r'\$(?!\s)', r'$ ', text)
            text = text.replace(placeholder, "$$")
            text = re.sub(r'(?<!\s)\$\$', r' $$', text)
            text = re.sub(r'\$\$(?!\s)', r'$$ ', text)
            return text

        answer = normalize_latex_spacing(answer)

        # 使用增强的图片路径匹配与规范化（仅在有检索结果时）
        if need_retrieval and original_image_paths:
            answer = normalize_markdown_images(answer, original_paths=original_image_paths)
        print("llm output", answer)

        # 生成建议回复问题
        suggested_replies = await generate_suggested_replies_baike(answer, extracted_text)

        return {
            "response": answer,
            "actions": [],
            "result_tag": "qa",
            "suggestedReplies": suggested_replies
        }
    except Exception as e:
        return {
            "response": f"生成回答失败：{str(e)}",
            "actions": [],
            "suggestedReplies": []
        }
def route_after_judge(state: QAState) -> str:
    """
    判断后的路由：根据need_retrieval决定下一步
    """
    need_retrieval = state.get("need_retrieval", True)
    if need_retrieval:
        return "rewrite_query_with_history"
    else:
        return "generate_answer"


# =============== 构建工作流 ===============
def create_baike_workflow():
    """创建普通问答工作流"""
    workflow = StateGraph(QAState)

    # 添加节点
    workflow.add_node("get_baike_history_context", get_baike_history_context)
    workflow.add_node("judge_need_retrieval", judge_need_retrieval)
    workflow.add_node("rewrite_query_with_history", rewrite_query_with_history)
    workflow.add_node("search_rag_and_graph", search_rag_and_graph)
    workflow.add_node("generate_answer", generate_answer)

    # 设置边
    workflow.add_edge(START, "get_baike_history_context")
    workflow.add_edge("get_baike_history_context", "judge_need_retrieval")
    workflow.add_conditional_edges(
        "judge_need_retrieval",
        route_after_judge,
        {
            "rewrite_query_with_history": "rewrite_query_with_history",
            "generate_answer": "generate_answer"
        }
    )
    workflow.add_edge("rewrite_query_with_history", "search_rag_and_graph")
    workflow.add_edge("search_rag_and_graph", "generate_answer")
    workflow.add_edge("generate_answer", END)

    # 编译工作流
    return workflow.compile()


# =============== 工作流执行函数 ===============
async def run_baike_workflow(
    extracted_text: str,
    combined_query: str,
    history_message: str = "",
    route_flag: str = ""
) -> Dict[str, Any]:
    """
    执行普通问答工作流（统一接口）

    Args:
        extracted_text: 文本描述（统一的输入参数）
        history_message: 历史消息（目前只接收，不做处理）

    Returns:
        工作流执行结果
    """
    # 创建并编译工作流
    app = create_baike_workflow()

    # 初始化状态（统一接口，只使用extracted_text和history_message）
    initial_state = {
        "extracted_text": extracted_text,
        "combined_query": combined_query,
        "retrieval_query": "",
        "route_flag": route_flag,
        "history_message": history_message,
        "history_messages": [],
        "need_retrieval": None,
        "rag_search_result": None,
        "graph_search_result": None,
        "response": "",
        "actions": [],
        "suggestedReplies": [],
        "error_message": ""
    }
    
    try:
        # 执行工作流（使用异步方式）
        final_state = await app.ainvoke(initial_state)
        
        # 检查是否有错误
        if final_state.get("error_message"):
            return {
                "response": f"普通问答工作流执行失败: {final_state['error_message']}",
                "actions": [],
                "result_tag": "qa",
                "suggestedReplies": []
            }
        
        # 统一输出格式：完全透传 generate_answer 的结果
        return {
            "response": final_state.get("response", ""),
            "actions": final_state.get("actions", []),
            "result_tag": "qa",
            "suggestedReplies": final_state.get("suggestedReplies", [])
        }
    
    except Exception as e:
        return {
            "response": f"普通问答工作流执行失败: {str(e)}",
            "actions": [],
            "result_tag": "qa",
            "suggestedReplies": []
        }



# =============== 测试 ===============
if __name__ == "__main__":
    import asyncio
    
    # 测试用例
    test_cases = [
        "船舰主发动机的安全警告是什么",
    ]
    
    async def test():
        for i, test_text in enumerate(test_cases, 1):
            print(f"\n{'='*60}")
            print(f"测试用例 {i}")
            print(f"{'='*60}")
            
            # ✅ 修复：传入 combined_query 参数
            result = await run_baike_workflow(
                extracted_text=test_text,
                combined_query=test_text
            )
            
            print(f"\n执行结果:")
            print(f"成功：{result.get('success', False)}")
            print(f"\n响应内容:")
            print(result.get("response", "无响应"))
    
    asyncio.run(test())