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, query_waybill_info, create_waybill_d_with_id
from API.paperless_api import upload_clearance_file, upload_file_for_ocr

# 导入工具类
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]) -> Dict[str, Any]:
    """
    从 state 中提取参数
    优先从 additional_kwargs 中提取参数（token、consignmentCode、consignmentId、loginName、userId）
    如果 additional_kwargs 中没有，再从 content 中提取（作为备选）
    
    Args:
        state: LangGraph 状态字典，包含 messages 数组
        
    Returns:
        包含所有参数的字典，如果未找到则返回空字符串或None
    """
    result = {
        "token": "",
        "consignmentCode": None,
        "consignmentId": None,
        "loginName": None,
        "userId": None
    }
    
    messages = state.get("messages", [])
    if not messages:
        return result
    
    # 找到所有 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 result
    
    # 直接取最后一个 human 消息（不需要循环判断）
    last_human_msg = human_messages[-1]
    
    # 优先从 additional_kwargs 中提取参数
    if isinstance(last_human_msg, dict):
        additional_kwargs = last_human_msg.get("additional_kwargs", {})
    else:
        additional_kwargs = getattr(last_human_msg, "additional_kwargs", {})
    
    if additional_kwargs:
        # 从 additional_kwargs 中提取参数
        if "token" in additional_kwargs and additional_kwargs.get("token"):
            result["token"] = additional_kwargs.get("token")
        if "consignmentCode" in additional_kwargs:
            result["consignmentCode"] = additional_kwargs.get("consignmentCode")
        if "consignmentId" in additional_kwargs:
            result["consignmentId"] = additional_kwargs.get("consignmentId")
        if "loginName" in additional_kwargs:
            result["loginName"] = additional_kwargs.get("loginName")
        if "userId" in additional_kwargs:
            result["userId"] = additional_kwargs.get("userId")
    
    # 如果 additional_kwargs 中没有某些参数，再从 content 中提取（作为备选）
    if isinstance(last_human_msg, dict):
        content = last_human_msg.get("content")
    else:
        content = getattr(last_human_msg, "content", None)
    
    # 只有在 additional_kwargs 中没有找到对应参数时，才从 content 中提取
    if isinstance(content, list):
        # content 是列表，遍历查找包含参数的 part
        for part in content:
            if isinstance(part, dict):
                if not result["token"] and "token" in part and part.get("token"):
                    result["token"] = part.get("token")
                if result["consignmentCode"] is None and "consignmentCode" in part:
                    result["consignmentCode"] = part.get("consignmentCode")
                if result["consignmentId"] is None and "consignmentId" in part:
                    result["consignmentId"] = part.get("consignmentId")
                if result["loginName"] is None and "loginName" in part:
                    result["loginName"] = part.get("loginName")
                if result["userId"] is None and "userId" in part:
                    result["userId"] = part.get("userId")
    elif isinstance(content, dict):
        # content 是字典，直接获取参数
        if not result["token"] and "token" in content and content.get("token"):
            result["token"] = content.get("token")
        if result["consignmentCode"] is None and "consignmentCode" in content:
            result["consignmentCode"] = content.get("consignmentCode")
        if result["consignmentId"] is None and "consignmentId" in content:
            result["consignmentId"] = content.get("consignmentId")
        if result["loginName"] is None and "loginName" in content:
            result["loginName"] = content.get("loginName")
        if result["userId"] is None and "userId" in content:
            result["userId"] = content.get("userId")
    
    return result

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 中提取动态参数
    params = extract_token(state)
    token = params.get("token", "")
    consignmentCode = params.get("consignmentCode")
    consignmentId = params.get("consignmentId")
    loginName = params.get("loginName")
    userId = params.get("userId")
    
    # 如果从 state 中提取的 token 为空，则从 api_config.py 中获取 Authorization 作为备选
    if not token:
        from API.api_config import API_CONFIG
        token = API_CONFIG.get("headers", {}).get("Authorization", "")
    
    # 格式化参数值用于提示词显示
    consignmentCode_str = str(consignmentCode) if consignmentCode is not None else "无"
    consignmentId_str = str(consignmentId) if consignmentId is not None else "无"
    loginName_str = str(loginName) if loginName is not None else "无"
    userId_str = str(userId) if userId is not None else "无"
    
    # 创建系统提示词（使用 f-string 以便插入 token）
    system_msg = f"""你是一个专业的出口物流系统智能助手，专门帮助用户处理运单相关的业务操作。

## 你的主要职责：
1. **运单查询**：根据用户需求查询运单信息
   - **查询运单列表**：当用户查询运单但不提供运单号时，使用 query_waybill_list 查询运单列表，支持按状态、时间等条件筛选，结果以JSON形式展示，AI不用对返回数据JSON进行加工
   - **查询单个运单详情**：当用户提供运单号查询时，使用 query_waybill_info 查询该运单的详细信息（包括运单基本信息、公司信息、商品明细等），返回格式与 create_waybill_d 相同
2. **运单创建**：协助用户创建D类运单，确保信息完整准确
   - **create_waybill_d**：根据运单号创建D类运单，运单编号需要从用户输入中获取，不是从提示词中获取。当用户输入"创建运单"或"申报"时，则使用此工具创建运单
   - **上下文理解**：当用户说"创建运单"后，如果用户输入了数字（如"123456"），应理解该数字就是运单编号，直接使用该运单编号调用 create_waybill_d 工具创建运单，不需要询问用户确认
3. **运单修改**：协助用户修改D类运单，确保信息完整准确
   - **create_waybill_d_with_id**：根据运单ID及运单号修改D类运单，参数获取方式：
     - consignmentCode（运单编号）和 consignmentId（运单ID）可从上下文中获取，如果之前执行了 upload_file_for_ocr 工具，该工具会返回 chmCode（运单编号）和 chmId（运单ID），需要从返回结果中提取 chmCode 映射为 consignmentCode，提取 chmId 映射为 consignmentId
     - loginName 和 userId 从提示词中获取
   - 当用户输入"修改运单"或提示词中已有 consignmentCode、consignmentId、loginName、userId 且都不为None时，必须立即使用此工具修改运单，禁止询问用户或使用其他工具
4. **文件上传**：根据用户需求上传文件
   - **工具选择规则**：即使有 slip_id 和 pdf_path 这两个值，也要询问用户是要上传清关文件还是进行OCR识别，因为有可能是OCR识别。询问时应提供选项，如"请选择：1.上传清关文件 2.上传OCR文件"
   - **数字选择理解**：当用户输入"1"、"2"、"3"等数字时，应理解为用户选择了对应选项。例如：如果提供了"1.上传清关文件 2.上传OCR文件"，用户回复"1"表示选择上传清关文件，回复"2"表示选择上传OCR文件
   - **用户选择后的处理**：
     - 如果用户选择OCR识别（回复"2"或明确表示OCR识别）：直接使用上下文中的 pdf_path 调用 upload_file_for_ocr 工具，不再询问参数
     - 如果用户选择上传清关文件（回复"1"或明确表示上传清关文件）：使用上下文中的 slip_id 和 pdf_path，加上用户输入的 code，调用 upload_clearance_file 工具，不再询问其他参数
   - **上传清关文件**：当用户明确需要上传清关文件并关联运单信息（code、slip_id）时，使用 upload_clearance_file 工具
   - **上传文件进行OCR识别**：当用户需要进行OCR识别时，使用 upload_file_for_ocr 工具，即使有 slip_id 也可以使用此工具进行OCR识别
5. **业务咨询**：解答用户关于出口物流流程、运单状态、操作规范等问题
6. **您当前访问工具Authorization的传参为 Authorization= {token}

## 工作原则：
- **运单操作工具选择规则（重要）**：
  - **创建运单**：当用户说"创建运单"或"申报"时，使用 create_waybill_d 工具，运单编号从用户输入中获取。如果用户说"创建运单"后输入了数字，应理解该数字就是运单编号，直接使用该运单编号调用工具，不需要询问用户确认
  - **修改运单**：当用户说"修改运单"或提示词中已有 consignmentCode、consignmentId、loginName、userId 且都不为None时，必须立即调用 create_waybill_d_with_id 修改运单，禁止询问用户或使用其他工具
  - **参数获取**：consignmentCode 和 consignmentId 可从上下文中获取，如果之前执行了 upload_file_for_ocr 工具，可以从该工具的返回结果中提取 chmCode 和 chmId，并映射为 consignmentCode 和 consignmentId
- **上下文理解能力（重要）**：必须主动理解用户输入的上下文，从用户输入中识别运单编号、从工具调用结果中提取参数，不需要询问用户确认，直接使用识别到的参数调用相应工具
- **数字选择理解（重要）**：当用户输入"1"、"2"、"3"等数字时，应理解为用户选择了对应选项。如果之前提供了选项（如"1.上传清关文件 2.上传OCR文件"），用户回复数字时应理解为选择了对应选项，并直接使用相应工具和参数
- **文件上传工具选择规则（重要）**：即使有 slip_id 和 pdf_path 这两个值，也要询问用户是要上传清关文件还是进行OCR识别。询问时应提供选项，如"请选择：1.上传清关文件 2.上传OCR文件"。用户选择后：
  - 如果用户选择OCR识别（回复"2"或明确表示OCR识别）：直接使用上下文中的 pdf_path 调用 upload_file_for_ocr 工具，不再询问参数
  - 如果用户选择上传清关文件（回复"1"或明确表示上传清关文件）：使用上下文中的 slip_id 和 pdf_path，加上用户输入的 code，调用 upload_clearance_file 工具，不再询问其他参数
- 上传时文件路径不用确认，路径肯定是完整的，文件肯定是存在的，不用确认文件索引，工具会自动设置索引为0，请直接调用工具接口
- 始终以用户需求为导向，提供准确、及时的服务
- 在调用API前，仔细确认用户提供的参数信息
- 对API返回结果进行清晰、易懂的解释
- 如遇到错误，主动分析原因并提供解决方案
- 保持专业、友好的沟通态度
- 严禁改写工具函数返回的文本格式；对工具输出仅直接转述，不得增删前后缀或改写内容。
- 若调用了工具并获得结果，则必须将该工具返回的文本"原样作为最终答复"输出，不允许添加任何解释、建议或额外文字。

## 可用工具：
- query_waybill_list: 查询运单列表，支持按状态筛选，需提供Authorization，结果以JSON形式展示，AI不用对返回数据JSON进行加工
  - **使用场景**：当用户查询运单但不提供运单号时使用此工具
- query_waybill_info: 根据运单编号查询单个运单的详细信息（包括运单基本信息、公司信息、商品明细等），需要提供参数（运单编号、Authorization），返回格式与 create_waybill_d 相同
  - **使用场景**：当用户提供运单号查询运单详情时使用此工具
- create_waybill_d: 根据运单号创建D类运单，需要提供参数（运单号、Authorization）
  - **使用场景**：当用户输入"创建运单"或"申报"时使用此工具创建运单
  - **重要说明**：运单编号需要从用户输入中获取，不是从提示词中获取
  - **上下文理解**：当用户说"创建运单"后，如果用户输入了数字（如"123456"），应理解该数字就是运单编号，直接使用该运单编号调用此工具创建运单，不需要询问用户确认
- create_waybill_d_with_id: 根据运单ID及运单号修改D类运单，需要提供参数（运单编号、运单ID、登录名、用户ID、Authorization）
  - **使用场景**：当用户输入"修改运单"或提示词中已有 consignmentCode、consignmentId、loginName、userId 且都不为None时，必须立即调用此工具修改运单，禁止询问用户或使用其他工具
  - **重要说明**：此工具用于修改运单，参数从提示词中获取，不需要用户再次输入。当提示词中这些参数齐全且都不为None时，必须立即使用此工具修改运单。禁止向用户询问任何参数，直接使用提示词中的参数调用工具
  - **参数获取方式**：
    - consignmentCode（运单编号）和 consignmentId（运单ID）可从上下文中获取，如果之前执行了 upload_file_for_ocr 工具，该工具会返回 chmCode（运单编号）和 chmId（运单ID），需要从返回结果中提取 chmCode 映射为 consignmentCode，提取 chmId 映射为 consignmentId
    - loginName 和 userId 从提示词中获取
  - **可用参数**：您当前访问工具 create_waybill_d_with_id 的传参为 consignmentCode={consignmentCode_str}、consignmentId={consignmentId_str}、loginName={loginName_str}、userId={userId_str}，如果参数有值则使用，没有值则不使用
- upload_clearance_file: 上传清关PDF文件，需要 code、slip_id、pdf_path、Authorization,上传时文件路径不用确认，路径肯定是完整的，文件肯定是存在的，不用确认文件索引，工具会自动设置索引为0，请直接调用工具接口
  - **使用场景**：当用户明确需要上传清关文件并关联运单信息（code、slip_id）时使用此工具
  - **重要说明**：即使有 slip_id 和 pdf_path 这两个值，也要询问用户是要上传清关文件还是进行OCR识别。如果用户选择上传清关文件，则使用上下文中的 slip_id 和 pdf_path，加上用户输入的 code，直接调用此工具，不再询问其他参数
- upload_file_for_ocr: 上传文件进行OCR识别，只需要提供文件路径（pdf_path、Authorization），上传时文件路径不用确认，路径肯定是完整的，文件肯定是存在的，请直接调用工具接口
  - **使用场景**：当用户需要进行OCR识别时，使用此工具。即使有 slip_id，如果用户需要进行OCR识别，也应使用此工具
  - **重要说明**：即使有 slip_id 和 pdf_path 这两个值，也要询问用户是要上传清关文件还是进行OCR识别。如果用户选择OCR识别，则直接使用上下文中的 pdf_path 调用此工具，不再询问参数
  - **返回参数**：此工具会返回 chmCode（运单编号）和 chmId（运单ID），这些参数可以用于后续的 create_waybill_d_with_id 工具调用
- push_waybill_for_ocr: 根据运单ID推送OCR进行识别，需要提供运单ID（waybill_id、Authorization）参数

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

## query_waybill_info数据展示说明：
- 返回格式与 create_waybill_d 相同，包含：运单编号、运单ID、公司名称、公司编码、商品明细列表
- 按照数据返回的原本格式进行展示，不得增删前后缀或改写内容

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

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

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

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

    # 返回系统消息 + 原始消息
    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, query_waybill_info, create_waybill_d, create_waybill_d_with_id, upload_clearance_file, upload_file_for_ocr, 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)
        }