import openai
import pytest
from openai.types.chat import ChatCompletionMessage

from agentdojo.agent_pipeline.llms.prompting_llm import PromptingLLM
from agentdojo.functions_runtime import FunctionCall
from agentdojo.types import ChatToolResultMessage, text_content_block_from_string


@pytest.fixture
def llm() -> PromptingLLM:
    client = openai.OpenAI(api_key="sk-1234")
    model = "llm-1234"
    return PromptingLLM(client, model)


def test_parse_model_output(llm: PromptingLLM):
    content = """\
Calling get_exchange_rate would be helpful because it can provide the exchange rate between USD and JPY, which is necessary to convert 1000 USD to JPY. However, the question lacks the required parameters for the function, specifically the base_currency and target_currency.

Fortunately, we can assume the base_currency is USD and the target_currency is JPY based on the context of the question.</function-thoughts>

<function-call>[get_exchange_rate(base_currency="USD", target_currency="JPY"), sum(a=1, b=2), f()]</function-call>

(Note: The actual function call and answer would depend on the implementation of the get_exchange_rate function, which is not provided here.)"""
    non_function_call_content = """\
Calling get_exchange_rate would be helpful because it can provide the exchange rate between USD and JPY, which is necessary to convert 1000 USD to JPY. However, the question lacks the required parameters for the function, specifically the base_currency and target_currency.

Fortunately, we can assume the base_currency is USD and the target_currency is JPY based on the context of the question.

(Note: The actual function call and answer would depend on the implementation of the get_exchange_rate function, which is not provided here.)"""

    message = ChatCompletionMessage(content=content, role="assistant")
    parsed = llm._parse_model_output(message)
    assert parsed["content"] is not None
    assert len(parsed["content"]) == 1
    assert parsed["content"][0]["type"] == "text"
    assert parsed["content"][0]["content"] == non_function_call_content

    assert parsed["tool_calls"] is not None
    assert len(parsed["tool_calls"]) == 3
    assert parsed["tool_calls"][0].function == "get_exchange_rate"
    assert parsed["tool_calls"][0].args == {
        "base_currency": "USD",
        "target_currency": "JPY",
    }
    assert parsed["tool_calls"][1].function == "sum"
    assert parsed["tool_calls"][1].args == {
        "a": 1,
        "b": 2,
    }
    assert parsed["tool_calls"][2].function == "f"
    assert parsed["tool_calls"][2].args == {}


def test_parse_model_output_no_call(llm: PromptingLLM):
    content = """\
<function-call>[]</function-call>
<answer>Hi, how are you?</answer>"""
    non_function_call_content = "Hi, how are you?"

    message = ChatCompletionMessage(content=content, role="assistant")
    parsed = llm._parse_model_output(message)
    assert parsed["content"] is not None
    assert len(parsed["content"]) == 1
    assert parsed["content"][0]["type"] == "text"
    assert parsed["content"][0]["content"] == non_function_call_content
    assert parsed["tool_calls"] is None or len(parsed["tool_calls"]) == 0


def test_tool_message_to_user_message_no_error(llm: PromptingLLM):
    tool_message = ChatToolResultMessage(
        tool_call=FunctionCall(
            function="get_exchange_rate",
            args={
                "base_currency": "USD",
                "target_currency": "JPY",
            },
        ),
        content=[text_content_block_from_string("100")],
        role="tool",
        tool_call_id=None,
        error=None,
    )
    user_message = llm._tool_message_to_user_message(tool_message)
    assert len(user_message["content"]) == 1
    assert user_message["content"][0]["type"] == "text"
    assert (
        user_message["content"][0]["content"]
        == """\
<function-call>get_exchange_rate(base_currency='USD', target_currency='JPY')</function-call><function-result>100</function-result>"""
    )


def test_tool_message_to_user_message_with_error(llm: PromptingLLM):
    tool_message = ChatToolResultMessage(
        tool_call=FunctionCall(
            function="get_exchange_rate",
            args={
                "base_currency": "USD",
                "target_currency": "JPY",
            },
        ),
        content=[text_content_block_from_string("")],
        role="tool",
        tool_call_id=None,
        error="ValueError: Invalid currency",
    )
    user_message = llm._tool_message_to_user_message(tool_message)
    assert len(user_message["content"]) == 1
    assert user_message["content"][0]["type"] == "text"
    assert (
        user_message["content"][0]["content"]
        == """\
<function-call>get_exchange_rate(base_currency='USD', target_currency='JPY')</function-call>\
<function-error>ValueError: Invalid currency</function-error>"""
    )
