from agentdojo.ast_utils import (
    create_python_function_from_tool_call,
    parse_tool_calls_from_python_function,
)
from agentdojo.functions_runtime import FunctionCall


def test_create_python_function_from_tool_call():
    tool_call = FunctionCall(
        function="get_exchange_rate",
        args={
            "base_currency": "USD",
            "target_currency": "JPY",
        },
    )
    expected = "get_exchange_rate(base_currency='USD', target_currency='JPY')"
    assert create_python_function_from_tool_call(tool_call) == expected


def test_parse_tool_calls_from_python_function():
    model_output = "[get_exchange_rate(base_currency='USD', target_currency='JPY'), sum(a=1, b=2), f()]"
    tool_calls = parse_tool_calls_from_python_function(model_output)
    assert len(tool_calls) == 3
    assert tool_calls[0]["function"] == "get_exchange_rate"
    assert tool_calls[0]["args"] == {"base_currency": "USD", "target_currency": "JPY"}
    assert tool_calls[1]["function"] == "sum"
    assert tool_calls[1]["args"] == {"a": 1, "b": 2}
    assert tool_calls[2]["function"] == "f"
    assert tool_calls[2]["args"] == {}


def test_parse_tool_calls_from_python_function_weird_args():
    model_output = "[send_email(recipients=['example@example.com'], subject='Hello', body='Hello', attachments={'type': 'document', 'id': '12'})]"
    tool_calls = parse_tool_calls_from_python_function(model_output)
    assert len(tool_calls) == 1
    assert tool_calls[0]["function"] == "send_email"
    assert tool_calls[0]["args"] == {
        "recipients": ["example@example.com"],
        "subject": "Hello",
        "body": "Hello",
        "attachments": {"type": "document", "id": "12"},
    }
