import json
import os
from openai import OpenAI


def get_current_weather(location: str, unit: str) -> str:
    """
    Get the current temperature at a location.

    Args:
        location: The location to get the temperature for, in the format "City, Country"
        unit: The unit to return the temperature in. (choices: ["celsius", "fahrenheit"])
    """

    return {
        "temperature": 20,
        "unit": unit,
    }

def calculate(expression: str) -> float:
    """
    Calculate the result of a mathematical expression.

    Args:
        expression: The mathematical expression to calculate. The expression must be a valid Python expression, e.g. "2 + 2" or "math.sqrt(16)".
    """

    return eval(expression, {}, {"math": __import__("math")})

def get_distance_between_cities(city1: str, city2: str) -> float:
    """
    Get the distance between two cities.

    Args:
        city1: The first city.
        city2: The second city.
    """

    return 100.0


# ---------------------------------------------------------------------------
# Tool schemas, in OpenAI function-calling format
# ---------------------------------------------------------------------------

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current temperature at a location.",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": 'The location to get the temperature for, in the format "City, Country".',
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "The unit to return the temperature in.",
                    },
                },
                "required": ["location", "unit"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Calculate the result of a mathematical expression.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "The mathematical expression to calculate. The expression must be a valid Python expression, e.g. '2 + 2' or 'math.sqrt(16)'.",
                    },
                },
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_distance_between_cities",
            "description": "Get the distance between two cities.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city1": {"type": "string", "description": "The first city."},
                    "city2": {"type": "string", "description": "The second city."},
                },
                "required": ["city1", "city2"],
            },
        },
    },
]

# Map tool name -> python callable
TOOL_MAP = {
    "get_current_weather": get_current_weather,
    "calculate": calculate,
    "get_distance_between_cities": get_distance_between_cities,
}


# ---------------------------------------------------------------------------
# Agent loop
# ---------------------------------------------------------------------------

BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.deepseek.com")
MODEL = os.environ.get("OPENAI_MODEL", "deepseek-v4-flash")
API_KEY = os.environ.get("OPENAI_API_KEY", "")
if not API_KEY:
    raise ValueError("OPENAI_API_KEY environment variable is not set.")

SYSTEM_PROMPT = (
    "You are a helpful assistant. Use the provided tools when they help "
    "answer the user's question, and always answer based on the tool results."
)

# create an OpenAI client that points to the llm server
client = OpenAI(base_url=BASE_URL, api_key=API_KEY)

def append_tool_calls_to_messages(messages, message):
    """Append the assistant message carrying the tool_calls array to the messages list.

    Per the OpenAI API, every role="tool" message must be preceded by an
    assistant message that declares the same tool_calls (with matching ids).
    """
    messages.append(
        {
            "role": "assistant",
            "content": message.content,
            "tool_calls": [
                {
                    "id": call.id,
                    "type": "function",
                    "function": {
                        "name": call.function.name,
                        "arguments": call.function.arguments,
                    },
                }
                for call in message.tool_calls
            ],
        }
    )

def loop(client: OpenAI) -> None:
    """Run a chat agent that can call tools, until the model gives a final answer."""

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
    ]
    while True:
        user_input = input(">>> ")
        if not user_input.strip():
            continue
        if user_input.lower() in ["exit", "quit"]:
            break
        
        messages.append({"role": "user", "content": user_input})
        resp = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=TOOLS,
        )
        while resp.choices[0].message.tool_calls:
            msg = resp.choices[0].message
            if msg.reasoning_content:
                print(f"\nAssistant (reasoning): {msg.reasoning_content}\n")

            # Model requested tool calls -> execute them and feed results back
            append_tool_calls_to_messages(messages, msg)

            for call in msg.tool_calls:
                fn_name = call.function.name
                try:
                    fn_args = json.loads(call.function.arguments or "{}")
                except json.JSONDecodeError:
                    fn_args = {}
                print("\ntool_call:", fn_name, fn_args)

                try:
                    result = TOOL_MAP[fn_name](**fn_args)
                except Exception as exc:
                    result = {"error": f"{type(exc).__name__}: {exc}"}

                print("tool_result:", result)
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call.id,
                        "content": json.dumps(result, ensure_ascii=False),
                    }
                )
            resp = client.chat.completions.create(
                model=MODEL,
                messages=messages,
                tools=TOOLS,
            )
        
        # No tool calls -> this is the final answer
        msg = resp.choices[0].message
        messages.append({"role": "assistant", "content": msg.content}) # add the final answer to the messages as chat history
        print(f"\nAssistant: {msg.content}\n")

if __name__ == "__main__":
    loop(client)