"""
主脑Agent - 简化版（无意图分类）

职责：
1. 输入处理：图片、音频、文本提取
2. 直接路由：根据外部传入的 route_flag 决定调用哪个工作流

工作流列表（来自 workflow_registry）：
- 故障排查与修理
- 舰船百科
- 操作使用
"""

from typing import TypedDict, Literal, Optional, Dict, Any, List
from langgraph.graph import StateGraph, START, END
from tools.function_tool import (
    vlm_image_to_text,
    asr_audio_to_text,
    detect_input_type
)
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls
from workflow_registry import VALID_ROUTE_FLAGS
import json
import re


class MainAgentState(TypedDict):
    """主脑Agent状态"""
    raw_input: Dict[str, Any]
    has_image: bool
    has_audio: bool
    has_query: bool
    has_file_text: bool
    input_type: str
    
    history_message: str
    extracted_text: str
    image_description: str
    asr_text: str
    file_text: str
    combined_query: str
    
    workflow_result: Optional[Dict[str, Any]]
    final_response: str
    error_message: str
    
    route_flag: str
    route_params: Dict[str, Any]


def detect_input_type_node(state: MainAgentState) -> Dict[str, Any]:
    """检测输入类型"""
    raw = state["raw_input"]
    result = detect_input_type.invoke({"raw_input": raw})
    
    has_image = result.get("has_image", False)
    has_audio = result.get("has_audio", False)
    has_query = result.get("has_query", False)
    
    has_file_text = False
    file_text_v = raw.get("file_text")
    if isinstance(file_text_v, str) and file_text_v.strip():
        has_file_text = True
    
    input_type = "mixed"
    count = sum([has_image, has_audio, has_query, has_file_text])
    if count == 1:
        if has_image:
            input_type = "image"
        elif has_audio:
            input_type = "audio"
        elif has_query:
            input_type = "text"
        elif has_file_text:
            input_type = "file"
    
    return {
        "has_image": has_image,
        "has_audio": has_audio,
        "has_query": has_query,
        "has_file_text": has_file_text,
        "input_type": input_type,
    }


async def process_multimodal_input(state: MainAgentState) -> Dict[str, Any]:
    """处理多模态输入"""
    has_image = state.get("has_image", False)
    has_audio = state.get("has_audio", False)
    has_query = state.get("has_query", False)
    has_file_text = state.get("has_file_text", False)
    raw = state["raw_input"]
    
    image_description = ""
    asr_text = ""
    extracted_text = ""
    
    try:
        if has_image:
            image_path = raw.get("image", "")
            if image_path:
                try:
                    image_description = await vlm_image_to_text.ainvoke({
                        "image_path": image_path
                    })
                    if not image_description or image_description.startswith("VLM 调用失败"):
                        image_description = ""
                except Exception:
                    image_description = ""
        
        if has_audio:
            audio_path = raw.get("audio", "")
            if audio_path:
                try:
                    asr_text = await asr_audio_to_text.ainvoke({
                        "audio_path": audio_path
                    })
                    if not asr_text or asr_text.startswith("ASR 调用失败"):
                        asr_text = ""
                except Exception:
                    asr_text = ""
        
        text_parts = []
        if raw.get("query"):
            text_parts.append(raw["query"])
        if image_description:
            text_parts.append(f"[图片描述] {image_description}")
        if asr_text:
            text_parts.append(f"[语音转文字] {asr_text}")
        if raw.get("file_text"):
            text_parts.append(f"[文件内容] {raw['file_text']}")
        
        extracted_text = "\n".join(text_parts) if text_parts else ""
        combined_query = raw.get("query", "") or extracted_text
        
        return {
            "extracted_text": extracted_text,
            "image_description": image_description,
            "asr_text": asr_text,
            "combined_query": combined_query,
            "file_text": raw.get("file_text", ""),
        }
    
    except Exception as e:
        return {
            "error_message": f"输入处理失败: {str(e)}",
            "extracted_text": "",
        }

# @track_function_calls
def route_to_workflow(state: MainAgentState) -> Dict[str, Any]:
    """
    🤖分析意图，调用相关智能体
    """
    route_flag = state.get("route_flag", "故障排查与修理")
    
    if route_flag not in VALID_ROUTE_FLAGS:
        route_flag = "故障排查与修理"
    
    query_text = state.get("extracted_text", "")
    combined_query = state.get("combined_query", "") or query_text
    
    raw_input = state.get("raw_input", {})
    history_message = raw_input.get("history_message", "")
    
    route_params = {
        "extracted_text": query_text,
        "combined_query": combined_query,
        "history_message": history_message,
    }
    
    if raw_input.get("file_text"):
        route_params["file_text"] = raw_input["file_text"]
    
    print(f"[主脑Agent] 路由到: {route_flag}")
    print(f"[主脑Agent] 参数: combined_query={combined_query[:100]}...")

    if route_flag == "故障排查与修理":
        detail = "正在调用故障排查与修理智能体"
    elif route_flag == "舰船百科":
        detail = "正在调用舰船百科智能体"
    elif route_flag == "操作使用":
        detail = "正在调用操作使用智能体"
    elif route_flag == "统计":
        detail = "正在调用统计智能体"
    else:
        detail = "正在调用通用智能体"
    
    return {
        "route_flag": route_flag,
        "route_params": route_params,
        "__detail__": detail,
    }


def create_main_agent():
    """创建主脑Agent"""
    workflow = StateGraph(MainAgentState)
    
    workflow.add_node("detect_input_type", detect_input_type_node)
    workflow.add_node("process_input", process_multimodal_input)
    workflow.add_node("route", route_to_workflow)
    
    workflow.add_edge(START, "detect_input_type")
    workflow.add_edge("detect_input_type", "process_input")
    workflow.add_edge("process_input", "route")
    workflow.add_edge("route", END)
    
    return workflow.compile()


from workflows.history_manager import filter_image_urls


async def extract_conversation_title(content: str) -> str:
    """提取对话标题"""
    try:
        prompt = f"请从以下对话内容中提取一个简短的标题（不超过20字）：\n\n{content[:500]}"
        response = await OpenaiAPI.open_api_chat_without_thinking(
            query=prompt,
            model=None,
            system_prompt="你是一个标题提取助手，请从对话内容中提取一个简短的标题。",
            messages=[]
        )
        title = response.strip() if response else "新对话"
        if len(title) > 30:
            title = title[:30] + "..."
        return title
    except Exception:
        return "新对话"


if __name__ == "__main__":
    import asyncio
    
    async def test():
        print("测试主脑Agent")
        agent = create_main_agent()
        result = await agent.ainvoke({
            "raw_input": {"query": "101舰主机发生异响"},
            "route_flag": "故障排查与修理"
        })
        print(f"结果: {result}")
    
    asyncio.run(test())

