from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
import os
import sys
import json
from typing import Dict, Any, List, Optional
from contextvars import ContextVar
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import AnyMessage

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

# 导入 API 模块
from API.waybill_api import query_waybill_list, create_waybill_d, push_waybill_for_ocr
from API.paperless_api import upload_clearance_file

# 导入工具类
from langgraph_examples.utils.message_processor import MessageProcessor

# 设置 DEEPSEEK API 配置
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-e59da2fbc73240ea8d5ef8fb12657e4b")
os.environ["OPENAI_BASE_URL"] = os.getenv("OPENAI_BASE_URL", "https://api.deepseek.com/v1")



# 创建 DEEPSEEK 聊天模型
model = ChatOpenAI(
    model="deepseek-chat",  # 使用 DEEPSEEK 模型
    temperature=0  # 固定输出，避免改写工具返回
)

## 直接传递函数作为工具

def pre_model_inspect_attachments(state, **kwargs):
    """
    LangGraph 预模型钩子：
    - 输入/输出都是"状态(dict)"，更新 'messages'
    - 发现文件/二进制分段：保存到目录，再把该段替换为纯文本 URL
    - 支持环境变量：
      ATTACH_SAVE_DIR  保存目录，默认 uploads
    """
    print("\n=== pre_model_hook: inspect attachments ===")
    try:
        messages = state.get("messages", [])

        # 结构化打印：处理前消息
        def _to_simple(msgs):
            out = []
            for m in msgs or []:
                if isinstance(m, dict):
                    out.append({"role": m.get("role"), "content": m.get("content")})
                else:
                    out.append({
                        "type": m.__class__.__name__,
                        "role": getattr(m, "role", None),
                        "content": getattr(m, "content", None),
                    })
            return out

        print("=== 处理前消息 ===")
        print(json.dumps(_to_simple(messages), ensure_ascii=False, indent=2))

        # 避免字符串与列表拼接导致异常，统一用结构化打印
        # print("处理前消息：", messages)
        
        # 使用工具类处理消息
        processor = MessageProcessor()
        filtered_messages, saved_files = processor.process_messages(messages)
        # 直接修改 state 中的 messages 结构，确保后续序列化使用新内容
        try:
            state["messages"] = filtered_messages
        except Exception:
            pass

        # 结构化打印：处理后消息
        print("=== 处理后消息 ===")
        print(json.dumps(_to_simple(filtered_messages), ensure_ascii=False, indent=2))

        if saved_files:
            print("=== saved files ===")
            for f in saved_files:
                print(f"  {f}")

        # 返回整个 state，避免上层忽略 messages 的替换
        return state
    except Exception as e:
        print(f"[pre_model_hook error] {e}")
        import traceback
        traceback.print_exc()
        return {}


def extract_token(state: Dict[str, Any]) -> str:
    """
    从 state 中提取 token
    获取最后一个类型为 HumanMessage 或 human 的消息中的 token
    
    Args:
        state: LangGraph 状态字典，包含 messages 数组
        
    Returns:
        token 字符串，如果未找到则返回空字符串
    """
    messages = state.get("messages", [])
    if not messages:
        return ""
    
    # 找到所有 is_human 类型的消息
    human_messages = []
    for msg in messages:
        # 兼容 dict 或 LangChain 的消息对象
        if isinstance(msg, dict):
            msg_type = msg.get("type")
        else:
            msg_type = msg.__class__.__name__
        
        # 检查是否是 human 类型的消息
        is_human = (msg_type == "HumanMessage" or msg_type == "human")
        if is_human:
            human_messages.append(msg)
    
    # 如果没有 human 消息，直接返回
    if not human_messages:
        return ""
    
    # 直接取最后一个 human 消息（不需要循环判断）
    last_human_msg = human_messages[-1]
    
    # 从 content 中提取 token
    if isinstance(last_human_msg, dict):
        content = last_human_msg.get("content")
    else:
        content = getattr(last_human_msg, "content", None)
    
    if isinstance(content, list):
        # content 是列表，遍历查找包含 token 的 part
        for part in content:
            if isinstance(part, dict) and "token" in part:
                token = part.get("token")
                if token:
                    return token
    elif isinstance(content, dict):
        # content 是字典，直接获取 token
        if "token" in content:
            token = content.get("token")
            if token:
                return token
    
    return ""

def _create_system_prompt(state: Dict[str, Any], config: RunnableConfig) -> List[AnyMessage]:
    """
    创建动态系统提示词
    
    Args:
        state: LangGraph 状态字典
        config: Runnable 配置
        
    Returns:
        包含系统消息和原始消息的列表
    """
    # 添加调试信息，确认函数被调用
    # print("\n=== _create_system_prompt 被调用 ===")
    # print(f"state type: {type(state)}")
    # print(f"state keys: {list(state.keys()) if isinstance(state, dict) else 'not a dict'}")
    
    # 从 state 中提取动态参数
    token = extract_token(state)
    
    # 如果从 state 中提取的 token 为空，则从 api_config.py 中获取 Authorization 作为备选
    if not token:
        from API.api_config import API_CONFIG
        token = API_CONFIG.get("headers", {}).get("Authorization", "")
    #
    # print(f"提取到的 token: {token[:30] if token else 'None'}...")
    
    # 创建系统提示词（使用 f-string 以便插入 token）
    system_msg = f"""你是一个专业的出口物流系统智能助手，专门帮助用户处理运单相关的业务操作。

## 你的主要职责：
1. **运单查询**：根据用户需求查询运单列表，支持按状态、时间等条件筛选，结果以JSON形式展示，AI不用对返回数据JSON进行加工
2. **运单创建**：协助用户创建D类运单，确保信息完整准确
   - **重要说明**：当用户输入"创建运单"或"申报"时，都理解为"创建运单"操作，应调用 create_waybill_d 工具
3. **业务咨询**：解答用户关于出口物流流程、运单状态、操作规范等问题
4. **您当前访问工具Authorization的传参为 Authorization= {token}

## 工作原则：
- 上传时文件路径不用确认，路径肯定是完整的，文件肯定是存在的，不用确认文件索引，工具会自动设置索引为0，请直接调用工具接口
- 始终以用户需求为导向，提供准确、及时的服务
- 在调用API前，仔细确认用户提供的参数信息
- 对API返回结果进行清晰、易懂的解释
- 如遇到错误，主动分析原因并提供解决方案
- 保持专业、友好的沟通态度
- 严禁改写工具函数返回的文本格式；对工具输出仅直接转述，不得增删前后缀或改写内容。
- 若调用了工具并获得结果，则必须将该工具返回的文本"原样作为最终答复"输出，不允许添加任何解释、建议或额外文字。

## 可用工具：
- query_waybill_list: 查询运单列表，支持按状态筛选，需提供Authorization，结果以JSON形式展示，AI不用对返回数据JSON进行加工
- create_waybill_d: 根据运单号创建D类运单，需要提供参数（运单号、Authorization）
  - **注意**：当用户说"创建运单"或"申报"时，都应调用此工具
- upload_clearance_file: 上传清关PDF文件，需要 code、slip_id、pdf_path、Authorization,上传时文件路径不用确认，路径肯定是完整的，文件肯定是存在的，不用确认文件索引，工具会自动设置索引为0，请直接调用工具接口
- push_waybill_for_ocr: 根据运单ID推送OCR进行识别，需要提供运单ID（waybill_id、Authorization）参数

## query_waybill_list数据展示说明：
- 运单查询结果会自动格式化为JSON形式展示，包含：运单号、运单类型、运单状态、发件人、运单日期
- 空字段会显示为空单元格
- 运单创建结果会显示成功/失败状态和详细信息

## create_waybill_d数据展示说明：
- 按照数据返回的原本格式进行展示，不得增删前后缀或改写内容。

请根据用户的具体需求，选择合适的工具并提供帮助。"""

    # 返回系统消息 + 原始消息
    result = [{"role": "system", "content": system_msg}] + state.get("messages", [])
    print(f"返回消息数量: {len(result)}")
    print("=== _create_system_prompt 执行完成 ===\n")
    return result


# 创建 ReAct 智能体
agent = create_react_agent(
    model=model,
    tools=[query_waybill_list, create_waybill_d, upload_clearance_file, push_waybill_for_ocr],
    pre_model_hook=pre_model_inspect_attachments,
    prompt=_create_system_prompt,
)

# 如果直接运行此文件
if __name__ == "__main__":
    # {"messages": [{"role": "user", "content": "查询状态为'等待录入'的运单列表"}]}
    # {"messages": [{"role": "user", "content": "帮我创建运单，运单编号：2025102904"}]}

    # 测试上传清关PDF文件（通过智能体调用 upload_clearance_file 工具）
    test_message = (
        "请调用工具 upload_clearance_file，并严格按以下参数执行：\n"
        "- code: 202510281\n"
        "- slip_id: 177950273\n"
        "- pdf_path: C:\\Users\\24790\\Desktop\\出口AI资料\\test2-1.pdf\n"
        "- file_index: 0\n"
        "- uid: 1761635727889\n"
        "只需执行工具并原样输出工具返回的文本，不要添加任何解释。"
    )
    result = agent.invoke({"messages": [{"role": "user", "content": test_message}]})
    print(result)

# LangGraph 服务端点
def api_agent_endpoint(input_data: Dict[str, Any]) -> Dict[str, Any]:
    """API 智能体服务端点"""
    try:
        result = agent.invoke(input_data)
        return {
            "status": "success",
            "data": result,
            "error": None
        }
    except Exception as e:
        return {
            "status": "error", 
            "data": None,
            "error": str(e)
        }