from unittest import mock

import pytest

from mlflow.genai.simulators.utils import (
    format_history,
    get_default_simulation_model,
    invoke_model_without_tracing,
)


@pytest.mark.parametrize(
    ("history", "expected"),
    [
        ([], None),
        ([{"role": "user", "content": "Hello"}], "user: Hello"),
        (
            [
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there!"},
                {"role": "user", "content": "How are you?"},
            ],
            "user: Hello\nassistant: Hi there!\nuser: How are you?",
        ),
        ([{"content": "Hello"}], "unknown: Hello"),
        ([{"role": "user"}], "user: "),
        ([{"role": None, "content": None}], "unknown: "),
    ],
)
def test_format_history(history, expected):
    assert format_history(history) == expected


@pytest.mark.parametrize(
    "model_uri",
    [
        "openai:/gpt-4o-mini",
        "anthropic:/claude-3-haiku",
    ],
)
def test_invoke_model_without_tracing_with_provider(model_uri):
    from mlflow.types.llm import ChatMessage

    messages = [ChatMessage(role="user", content="Hello")]

    with mock.patch(
        "mlflow.genai.scorers.llm_backend.ScorerLLMClient.complete", return_value="Hi there!"
    ) as mock_invoke:
        result = invoke_model_without_tracing(model_uri=model_uri, messages=messages)

        assert result == "Hi there!"
        mock_invoke.assert_called_once()


def test_invoke_model_without_tracing_with_inference_params():
    from mlflow.types.llm import ChatMessage

    messages = [ChatMessage(role="user", content="Hello")]

    with mock.patch(
        "mlflow.genai.scorers.llm_backend.ScorerLLMClient.complete", return_value="Response"
    ) as mock_invoke:
        invoke_model_without_tracing(
            model_uri="openai:/gpt-4o-mini",
            messages=messages,
            inference_params={"temperature": 0.5},
        )

        mock_invoke.assert_called_once_with(
            [{"role": "user", "content": "Hello"}],
            response_format=None,
            num_retries=3,
            temperature=0.5,
        )


@pytest.mark.parametrize("model_uri", ["databricks", "gpt-oss-120b"])
def test_invoke_model_without_tracing_with_databricks(model_uri):
    from mlflow.types.llm import ChatMessage

    messages = [ChatMessage(role="user", content="Hello")]

    with (
        mock.patch("mlflow.genai.simulators.utils.call_chat_completions") as mock_call,
        mock.patch(
            "mlflow.genai.simulators.utils._create_message_from_databricks_response"
        ) as mock_create,
    ):
        mock_call.return_value = mock.MagicMock(error_code=None, output_json='{"content": "Hi"}')
        mock_create.return_value = mock.MagicMock(content="Hi from Databricks")

        result = invoke_model_without_tracing(model_uri=model_uri, messages=messages)

        assert result == "Hi from Databricks"
        mock_call.assert_called_once()


def test_get_default_simulation_model_non_databricks():
    with mock.patch("mlflow.genai.simulators.utils.is_databricks_uri", return_value=False):
        model = get_default_simulation_model()
        assert model == "openai:/gpt-5"


def test_get_default_simulation_model_databricks():
    with mock.patch("mlflow.genai.simulators.utils.is_databricks_uri", return_value=True):
        model = get_default_simulation_model()
        assert model == "gpt-oss-120b"