import json

import pytest
from langchain.prompts import ChatPromptTemplate
from langchain.schema.exceptions import LangChainException
from langchain.schema.messages import HumanMessage

from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError

pytest_plugins = ("pytest_asyncio",)


def test_bedrock_caller_load_settings(mocker, mock_settings):
    mocked_boto3_client = mocker.patch(
        "llm_api.backends.bedrock.BedrockCaller.get_boto3_client"
    )
    mocked_bedrock_client = mocker.patch(
        "llm_api.backends.bedrock.BedrockCaller.get_client"
    )

    caller = BedrockCaller(mock_settings)
    expected_test_key = mock_settings.aws_secret_access_key.get_secret_value()

    assert caller.settings.aws_secret_access_key.get_secret_value() == expected_test_key
    mocked_boto3_client.assert_called_once()
    mocked_bedrock_client.assert_called_once()


def test_generate_openai_prompt_success():
    user_input = "What day is it today?"

    prompt_output = BedrockCaller.generate_prompt()
    assert isinstance(prompt_output, ChatPromptTemplate)

    prompt_output = prompt_output.format_messages(text=user_input)

    expected_prompt_elements = 5
    assert len(prompt_output) == expected_prompt_elements
    assert isinstance(prompt_output[-1], HumanMessage)
    assert prompt_output[-1].content == user_input


@pytest.mark.asyncio
async def test_call_model_success(mocker, mock_settings):
    caller = BedrockCaller(mock_settings)

    expected_entities = ["William Shakespeare", "Globe Theatre"]
    mocked_result = {
        "entities": [
            {
                "uri": "William Shakespeare",
                "description": "English playwright, poet, and actor",
                "wikipedia_url": "https://en.wikipedia.org/wiki/William_Shakespeare",
            },
            {
                "uri": "Globe Theatre",
                "description": "Theatre in London associated with William Shakespeare",
                "wikipedia_url": "https://en.wikipedia.org/wiki/Globe_Theatre",
            },
        ],
        "connections": [
            {
                "from": "William Shakespeare",
                "to": "Globe Theatre",
                "label": "performed plays at",
            },
        ],
        "user_search": "Who is Shakespeare?",
    }
    mocked_result = "Test json ```json" + json.dumps(mocked_result) + "```"
    mocker.patch(
        "langchain.schema.runnable.base.RunnableSequence.ainvoke",
        return_value=mocked_result,
    )
    user_template = "{text}"
    test_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a test system"),
            ("system", "Provide a valid JSON response to the user."),
            ("user", user_template),
        ]
    )

    test_search = "Who is Shakespeare?"
    response = await caller.call_model(test_prompt, test_search)

    assert expected_entities == [entity["uri"] for entity in response["entities"]]


@pytest.mark.asyncio
async def test_call_model_failure_index_error(mocker, mock_settings):
    caller = BedrockCaller(mock_settings)
    mocked_result = {
        "entities": [
            {
                "uri": "William Shakespeare",
                "description": "English playwright, poet, and actor",
                "wikipedia_url": "https://en.wikipedia.org/wiki/William_Shakespeare",
            },
            {
                "uri": "Globe Theatre",
                "description": "Theatre in London associated with William Shakespeare",
                "wikipedia_url": "https://en.wikipedia.org/wiki/Globe_Theatre",
            },
        ],
        "connections": [
            {
                "from": "William Shakespeare",
                "to": "Globe Theatre",
                "label": "performed plays at",
            },
        ],
        "user_search": "Who is Shakespeare?",
    }
    mocked_result = "Test json" + json.dumps(mocked_result)
    mocker.patch(
        "langchain.schema.runnable.base.RunnableSequence.ainvoke",
        return_value=mocked_result,
    )
    expected_error_message = "Unable to parse model output as expected."

    user_template = "{text}"
    test_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a test system"),
            ("system", "Provide a valid JSON response to the user."),
            ("user", user_template),
        ]
    )
    test_search = "Who is Shakespeare?"

    with pytest.raises(BedrockModelCallError) as exception:
        await caller.call_model(test_prompt, test_search)
    assert expected_error_message in str(exception.value)


@pytest.mark.asyncio
async def test_call_model_failure_json_decode_error(mocker, mock_settings):
    caller = BedrockCaller(mock_settings)
    mocked_result = {
        "entities": [
            {
                "uri": "William Shakespeare",
                "description": "English playwright, poet, and actor",
                "wikipedia_url": "https://en.wikipedia.org/wiki/William_Shakespeare",
            },
            {
                "uri": "Globe Theatre",
                "description": "Theatre in London associated with William Shakespeare",
                "wikipedia_url": "https://en.wikipedia.org/wiki/Globe_Theatre",
            },
        ],
        "connections": [
            {
                "from": "William Shakespeare",
                "to": "Globe Theatre",
                "label": "performed plays at",
            },
        ],
        "user_search": "Who is Shakespeare?",
    }
    mocked_result = "Test ```json," + json.dumps(mocked_result) + "```"
    mocker.patch(
        "langchain.schema.runnable.base.RunnableSequence.ainvoke",
        return_value=mocked_result,
    )
    expected_error_message = "Error decoding model output."

    user_template = "{text}"
    test_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a test system"),
            ("system", "Provide a valid JSON response to the user."),
            ("user", user_template),
        ]
    )
    test_search = "Who is Shakespeare?"

    with pytest.raises(BedrockModelCallError) as exception:
        await caller.call_model(test_prompt, test_search)
    assert expected_error_message in str(exception.value)


@pytest.mark.asyncio
async def test_call_model_failure_api_error(mocker, mock_settings):
    caller = BedrockCaller(mock_settings)

    mocked_client_call = mocker.patch(
        "langchain.schema.runnable.base.RunnableSequence.ainvoke"
    )
    expected_error_message = "Error calling model."

    mocked_client_call.side_effect = ValueError()

    user_template = "{text}"
    test_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a test system"),
            ("system", "Provide a valid JSON response to the user."),
            ("user", user_template),
        ]
    )
    test_search = "Who is Shakespeare?"

    with pytest.raises(BedrockModelCallError) as exception:
        await caller.call_model(test_prompt, test_search)

    assert expected_error_message in str(exception.value)


@pytest.mark.asyncio
async def test_call_model_failure_langchain_error(mocker, mock_settings):
    caller = BedrockCaller(mock_settings)

    mocked_client_call = mocker.patch(
        "langchain.schema.runnable.base.RunnableSequence.ainvoke"
    )
    expected_error_message = "Error sending prompt to LLM."

    mocked_client_call.side_effect = LangChainException()

    user_template = "{text}"
    test_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a test system"),
            ("system", "Provide a valid JSON response to the user."),
            ("user", user_template),
        ]
    )
    test_search = "Who is Shakespeare?"

    with pytest.raises(BedrockModelCallError) as exception:
        await caller.call_model(test_prompt, test_search)

    assert expected_error_message in str(exception.value)

langchain.prompts.ChatPromptTemplate.from_messages langchain.schema.exceptions.LangChainException
[((597, 625), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (610, 625), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((994, 1025), 'llm_api.backends.bedrock.BedrockCaller.generate_prompt', 'BedrockCaller.generate_prompt', ([], {}), '()\n', (1023, 1025), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((1443, 1471), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (1456, 1471), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((2561, 2724), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (["[('system', 'You are a test system'), ('system',\n    'Provide a valid JSON response to the user.'), ('user', user_template)]"], {}), "([('system', 'You are a test system'), (\n    'system', 'Provide a valid JSON response to the user.'), ('user',\n    user_template)])\n", (2593, 2724), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((3073, 3101), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (3086, 3101), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((4183, 4346), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (["[('system', 'You are a test system'), ('system',\n    'Provide a valid JSON response to the user.'), ('user', user_template)]"], {}), "([('system', 'You are a test system'), (\n    'system', 'Provide a valid JSON response to the user.'), ('user',\n    user_template)])\n", (4215, 4346), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((4728, 4756), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (4741, 4756), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((5837, 6000), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (["[('system', 'You are a test system'), ('system',\n    'Provide a valid JSON response to the user.'), ('user', user_template)]"], {}), "([('system', 'You are a test system'), (\n    'system', 'Provide a valid JSON response to the user.'), ('user',\n    user_template)])\n", (5869, 6000), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((6374, 6402), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (6387, 6402), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((6666, 6829), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (["[('system', 'You are a test system'), ('system',\n    'Provide a valid JSON response to the user.'), ('user', user_template)]"], {}), "([('system', 'You are a test system'), (\n    'system', 'Provide a valid JSON response to the user.'), ('user',\n    user_template)])\n", (6698, 6829), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((7210, 7238), 'llm_api.backends.bedrock.BedrockCaller', 'BedrockCaller', (['mock_settings'], {}), '(mock_settings)\n', (7223, 7238), False, 'from llm_api.backends.bedrock import BedrockCaller, BedrockModelCallError\n'), ((7449, 7469), 'langchain.schema.exceptions.LangChainException', 'LangChainException', ([], {}), '()\n', (7467, 7469), False, 'from langchain.schema.exceptions import LangChainException\n'), ((7518, 7681), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (["[('system', 'You are a test system'), ('system',\n    'Provide a valid JSON response to the user.'), ('user', user_template)]"], {}), "([('system', 'You are a test system'), (\n    'system', 'Provide a valid JSON response to the user.'), ('user',\n    user_template)])\n", (7550, 7681), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((3909, 3934), 'json.dumps', 'json.dumps', (['mocked_result'], {}), '(mocked_result)\n', (3919, 3934), False, 'import json\n'), ((4449, 4485), 'pytest.raises', 'pytest.raises', (['BedrockModelCallError'], {}), '(BedrockModelCallError)\n', (4462, 4485), False, 'import pytest\n'), ((6103, 6139), 'pytest.raises', 'pytest.raises', (['BedrockModelCallError'], {}), '(BedrockModelCallError)\n', (6116, 6139), False, 'import pytest\n'), ((6932, 6968), 'pytest.raises', 'pytest.raises', (['BedrockModelCallError'], {}), '(BedrockModelCallError)\n', (6945, 6968), False, 'import pytest\n'), ((7784, 7820), 'pytest.raises', 'pytest.raises', (['BedrockModelCallError'], {}), '(BedrockModelCallError)\n', (7797, 7820), False, 'import pytest\n'), ((2353, 2378), 'json.dumps', 'json.dumps', (['mocked_result'], {}), '(mocked_result)\n', (2363, 2378), False, 'import json\n'), ((5568, 5593), 'json.dumps', 'json.dumps', (['mocked_result'], {}), '(mocked_result)\n', (5578, 5593), False, 'import json\n')]