"""
工作流：普通问答Agent（使用LangGraph）
输入：文本（转化后和原文本query）、图片描述、语音文本
步骤：
1. 使用RAG和GraphRAG检索相关信息
2. 生成回答
使用 OpenaiAPI 方法调用大模型，尽可能使用异步
"""
from typing import TypedDict, Optional, Dict, Any, List
from langgraph.graph import StateGraph, START, END
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls, get_current_stream_handler
import asyncio
import json
import re


# =============== 定义状态 ===============
class Othertate(TypedDict):
    """普通问答工作流状态"""
    extracted_text: str  # 已提取的文本（统一接口，只接收文本）
    combined_query: str
    route_flag: str  # 路由标识
    history_message: str  # 历史消息（JSON格式）
    history_messages: List[Dict[str, Any]]  # 过滤后的历史消息列表

    # 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_other_history_context(state: Othertate) -> Dict[str, Any]:
    """
    知识问答步骤：获取问答类历史信息
    问答工作流只关注历史问答记录，用于提供上下文
    支持格式：[{'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'}]
    返回消息列表（history_message已经在app.py中过滤过了）
    """
    import json

    history_message = state.get("history_message", "") or ""
    history_messages = []

    if history_message:
        try:
            # history_message已经在app.py中过滤过了，直接解析
            if isinstance(history_message, str):
                history_data = json.loads(history_message)
            else:
                history_data = history_message
            
            if isinstance(history_data, list):
                # 只取最近10条历史记录，避免上下文过长
                history_messages = history_data[-10:]
        except Exception as e:
            pass
    
    return {"history_messages": history_messages}


async def generate_suggested_replies_other(answer: str, question: str) -> List[Dict[str, Any]]:
    """
    生成针对问答结果的建议回复问题
    使用大模型生成两个相关问题
    """
    prompt = f"""
你是一个问答助手。根据以下问答内容，生成两个用户可能关心的后续问题。

【用户问题】
{question}

【回答内容】
{answer[:500]}

请生成两个简洁、实用的问题，这些问题应该：
1. 不得出现需要归档和下载等操作相关的问题
2. 与当前问答内容相关，能够帮助用户进一步了解相关信息
3. 问题要具体、可操作
4. 每个问题不超过20个字

请以JSON格式输出，格式如下：
[
    {{"title": "问题1的标题", "content": "问题1的完整内容"}},
    {{"title": "问题2的标题", "content": "问题2的完整内容"}}
]

只输出JSON，不要有其他文字说明。
"""
    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(prompt, model=None)
        # 尝试解析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: Othertate) -> Dict[str, Any]:
    """
    正在生成问答结果
    """
    extracted_text = state.get("combined_query", "")
    history_messages = state.get("history_messages", []) or []
    print("history_message(qa)", history_messages)

    try:
        # 构建系统提示词
        system_prompt = """我是AI助手，根据历史信息和用户问题进行回答。
请基于以上信息，给出准确、专业的回答。如果有历史问答上下文，请结合历史对话内容理解当前问题。"""

        # 构建消息列表
        messages = []
        
        # 添加历史消息
        if history_messages:
            messages.extend(history_messages)
        
        # 添加当前用户查询
        messages.append({"role": "user", "content": f"用户问题：{extracted_text}"})

        # 调用LLM生成回答（异步，禁用thinking）
        answer = await OpenaiAPI.open_api_chat_without_thinking(
            model=None,
            system_prompt=system_prompt,
            messages=messages
        )

        # 过滤掉 <think> 标签（如果有）
        import re
        answer = re.sub(r'<think>.*?</think>', '', answer, flags=re.DOTALL)
        answer = answer.strip() if answer else "抱歉，无法生成回答。"

        # 返回按钮：确认结果、重新生成
        actions = []

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

        return {
            "response": answer,
            "actions": actions,
            "result_tag": "other",
            "suggestedReplies": suggested_replies
        }
    except Exception as e:
        return {
            "response": f"生成回答失败: {str(e)}",
            "actions": [],
            "suggestedReplies": []
        }


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

    # 添加节点
    workflow.add_node("get_other_history_context", get_other_history_context)
    workflow.add_node("generate_answer", generate_answer)

    # 设置边
    workflow.add_edge(START, "get_other_history_context")
    workflow.add_edge("get_other_history_context", "generate_answer")
    workflow.add_edge("generate_answer", END)

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


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

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

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

    # 初始化状态（统一接口，只使用extracted_text和history_message）
    initial_state = {
        "extracted_text": extracted_text,
        "combined_query": combined_query,
        "history_message": history_message,
        "history_messages": [],
        "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": "other",
            "suggestedReplies": final_state.get("suggestedReplies", [])
        }

    except Exception as e:
        return {
            "response": f"普通问答工作流执行失败: {str(e)}",
            "actions": [],
            "result_tag": "other",
            "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}")

            result = await run_other_workflow(extracted_text=test_text)

            print(f"\n执行结果:")
            print(f"成功: {result.get('success', False)}")
            print(f"\n响应内容:")
            print(result.get("response", "无响应"))


    asyncio.run(test())