from langchain.chat_models import init_chat_model
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import before_model, wrap_tool_call, AgentState

# 使用 DeepSeek 模型，通过阿里云百炼平台接入
# init_chat_model 会自动适配 OpenAI 兼容接口
model = init_chat_model(
    model="deepseek-v4-flash",
    model_provider="openai",
    base_url="https://api.deepseek.com",
    api_key="sk-80a123483afb480285c6452985eea18e",
)
# ---- 模拟数据 ----
# 实际项目中，这些数据来自数据库或 API
EXPORT_RIGHTS = {
    "zhangsan": {"role": "finance", "region": "all"},
    "lisi": {"role": "operation", "region": "east"},
}


@tool
def check_export_permission(username: str) -> dict:
    """查询用户是否有报表导出权限。
    参数 username 为员工账号（英文名）。
    """
    user_info = EXPORT_RIGHTS.get(username)
    if not user_info:
        return {"username": username, "can_export": False, "reason": "用户不存在"}
    return {
        "username": username,
        "role": user_info["role"],
        "region": user_info["region"],
        "can_export": user_info["role"] in ("finance", "ops_manager"),
    }


@tool
def create_export_task(
    report_name: str, file_format: str, estimated_rows: int, reason: str
) -> dict:
    """创建报表导出任务。
    report_name: 报表名称
    file_format: 导出格式 (xlsx / csv)
    estimated_rows: 预计导出行数
    reason: 导出原因
    """
    return {
        "task_id": "EXPORT-20260706-001",
        "report_name": report_name,
        "file_format": file_format,
        "estimated_rows": estimated_rows,
        "reason": reason,
        "status": "queued",
    }


tools = [check_export_permission, create_export_task]



@before_model
def log_before_model(state: AgentState, runtime):
    """模型调用前执行：看一眼当前上下文里有多少条消息。"""
    # state["messages"] 是 Agent 当前累积的全部对话历史
    # runtime 携带运行时环境信息（后面会详细讲）
    print(f"[审计] 准备调用模型，当前上下文消息数：{len(state['messages'])}")
    return None  # 返回 None 表示不做任何修改


@wrap_tool_call
def log_tool_call(request, handler):
    """工具调用前后各打一条日志，并记录耗时。"""
    tool_name = request.tool_call["name"]
    print(f"[审计] 开始执行工具：{tool_name}")

    # handler(request) 是"真正执行工具"的入口
    # 不调用它，工具就不会执行
    result = handler(request)

    print(f"[审计] 工具执行完毕：{tool_name}")
    return result


# 组装 Agent，把 Middleware 列表传进去
agent = create_agent(
    model=model,
    tools=tools,
    middleware=[
        log_before_model,   # 排在前面：先记录状态
        log_tool_call,      # 排在后面：包裹工具调用
    ],
    system_prompt=(
        "你是报表导出平台的智能助手。"
        "用户询问导出相关问题前，先调用 check_export_permission 确认权限。"
        "只有在用户明确提出导出需求时，才调用 create_export_task 创建任务。"
    ),
)

# 跑一次看看
response = agent.invoke({
    "messages": [
        {
            "role": "user",
            "content": "我是 zhangsan，需要导出本月 east 区域的订单报表，大约 5000 行，xlsx 格式。",
        }
    ]
})

print(response["messages"][-1].content)