import os
from langchain.agents import AgentType, initialize_agent
from langchain.tools import StructuredTool
from langchain_deepseek import ChatDeepSeek
from pydantic import BaseModel, Field
from typing import Optional
import math
import requests

class CalculatorInput(BaseModel):
    operation: str = Field(description="Mathematical operation: add, subtract, multiply, divide, power, sqrt")
    a: float = Field(description="First number")
    b: Optional[float] = Field(default=None, description="Second number (not required for sqrt)")

def calculator(operation: str, a: float, b: Optional[float] = None) -> str:
    """Perform basic arithmetic operations."""
    op = operation.lower()
    
    if op == "add":
        return f"{a} + {b} = {a + b}"
    elif op == "subtract":
        return f"{a} - {b} = {a - b}"
    elif op == "multiply":
        return f"{a} * {b} = {a * b}"
    elif op == "divide":
        if b == 0:
            return "Error: Division by zero"
        return f"{a} / {b} = {a / b}"
    elif op == "power":
        return f"{a}^{b} = {a ** b}"
    elif op == "sqrt":
        if a < 0:
            return "Error: Cannot take square root of negative number"
        return f"sqrt({a}) = {math.sqrt(a)}"
    else:
        return f"Unknown operation: {op}"

calculator_tool = StructuredTool.from_function(
    func=calculator,
    name="Calculator",
    description="Useful for performing mathematical calculations. Input should include operation (add, subtract, multiply, divide, power, sqrt) and numbers.",
    args_schema=CalculatorInput
)

class WeatherInput(BaseModel):
    city: str = Field(description="Name of the city to get the weather for")

def get_weather(city: str) -> str:
    """模拟天气信息（不访问真实 API）"""
    mock_data = {
        "北京": {"description": "晴天", "temp": 26},
        "上海": {"description": "小雨", "temp": 22},
        "广州": {"description": "多云", "temp": 30},
        "杭州": {"description": "雷阵雨", "temp": 24}
    }

    weather = mock_data.get(city, {"description": "未知", "temp": "N/A"})
    return f"{city} 当前天气：{weather['description']}，温度：{weather['temp']}°C"

weather_tool = StructuredTool.from_function(
    func=get_weather,
    name="WeatherTool",
    description="查询指定城市的天气信息（模拟数据）",
    args_schema=WeatherInput
)


llm = ChatDeepSeek(api_key=os.getenv("DEEPSEEK_API_KEY"), model="deepseek-chat")


tools = [calculator_tool, weather_tool]

agent = initialize_agent(
    tools, 
    llm, 
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION, 
    verbose=True
)

if __name__ == "__main__":
    while True:
        query = input("\nEnter your question (or 'exit' to quit): ")
        if query.lower() == 'exit':
            break
        result = agent.invoke(query)
        print(f"Result: {result['output']}")