from langchain_openai import ChatOpenAI
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_core.tools import StructuredTool, Tool
from langchain.agents import create_agent
import asyncio

def normalize_base_url(url):
    """标准化base_url，移除反引号和首尾空格，去除末尾斜杠"""
    if url:
        url = url.strip().strip('`').rstrip('/')
    return url

# OpenAI配置
api_key = "sk-8d657b8b7efe0cb6c141a30d9cee97f726efb9b18ea72bb1e8cfb080b42c140d"
base_url = normalize_base_url("https://console.pivotbak.cfd/v1")
model_name = "MiniMax-M2.7-highspeed"

# MCP配置
mcp_server_url = normalize_base_url("http://192.168.0.10:9999/admin/mcp/sse")

# 全局参数 - 每次对话都要传递
team_id = 19
token = "f1d23a59-479c-44a8-94b5-2344e4672ffd"

# 初始化LLM，每次请求都会携带teamId和token
llm = ChatOpenAI(
    model=model_name,
    openai_api_key=api_key,
    openai_api_base=base_url,
    temperature=0.7,
    max_tokens=4096,
    extra_body={
        "teamId": team_id,
        "token": token
    }
)

# 示例工具函数
def get_current_time():
    """获取当前时间"""
    from datetime import datetime
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def calculate(a: float, b: float, operation: str = "add") -> float:
    """
    简单计算器
    :param a: 第一个数
    :param b: 第二个数
    :param operation: 操作类型，可选 add, sub, mul, div
    """
    if operation == "add":
        return a + b
    elif operation == "sub":
        return a - b
    elif operation == "mul":
        return a * b
    elif operation == "div":
        if b == 0:
            return "错误：除数不能为零"
        return a / b
    else:
        return f"未知操作: {operation}"

# 定义工具
tools = [
    Tool(
        name="get_current_time",
        func=get_current_time,
        description="获取当前时间"
    ),
    StructuredTool.from_function(calculate)
]

# 从MCP加载工具
async def load_mcp_tools():
    mcp_client = MultiServerMCPClient(
        {"server": {"url": mcp_server_url, "transport": "sse"}}
    )
    try:
        mcp_tools = await mcp_client.get_tools()
        print(f"从MCP服务器加载到 {len(mcp_tools)} 个工具:")
        for tool in mcp_tools:
            print(f"  - {tool.name}: {tool.description}")
        return mcp_tools
    except Exception as e:
        print(f"连接MCP服务器失败: {e}")
        print("将继续使用本地工具")
        return []

# 创建Agent
def create_my_agent(all_tools):
    agent = create_agent(
        model=llm,
        tools=all_tools,
        system_prompt="你是一个有用的助手，使用提供的工具来回答问题。"
    )
    return agent

# 测试示例
async def main():
    print("=" * 60)
    print("LangChain + MCP 测试")
    print(f"模型: {model_name}")
    print(f"API地址: {base_url}")
    print(f"MCP地址: {mcp_server_url}")
    print(f"Team ID: {team_id}")
    print(f"Token: {token[:10]}..." if len(token) > 10 else token)
    print("=" * 60)

    # 加载MCP工具
    mcp_tools = await load_mcp_tools()
    all_tools = tools + mcp_tools

    # 创建Agent
    agent = create_my_agent(all_tools)

    # 测试1: 直接对话
    print("\n--- 测试1: 直接对话 ---")
    response = llm.invoke("你好，介绍一下自己")
    print(response.content)

    # 测试2: 使用工具
    print("\n--- 测试2: 使用工具 - 获取当前时间 ---")
    result = await agent.ainvoke({
        "messages": [("user", "现在几点了？")]
    })
    print(result["messages"][-1].content)

    # 测试3: 使用计算器
    print("\n--- 测试3: 使用工具 - 计算 ---")
    result = await agent.ainvoke({
        "messages": [("user", "计算 100 乘以 50 等于多少？")]
    })
    print(result["messages"][-1].content)

if __name__ == "__main__":
    asyncio.run(main())