#  Copyright (c) "Neo4j"
#  Neo4j Sweden AB [https://neo4j.com]
#  #
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  #
#      https://www.apache.org/licenses/LICENSE-2.0
#  #
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from unittest import mock
from unittest.mock import MagicMock, call

import pytest
from neo4j_graphrag.exceptions import RagInitializationError, SearchValidationError
from neo4j_graphrag.generation.graphrag import GraphRAG
from neo4j_graphrag.generation.prompts import RagTemplate
from neo4j_graphrag.generation.types import RagResultModel
from neo4j_graphrag.llm import LLMResponse
from neo4j_graphrag.message_history import InMemoryMessageHistory
from neo4j_graphrag.types import LLMMessage, RetrieverResult, RetrieverResultItem


def test_graphrag_prompt_template() -> None:
    template = RagTemplate()
    prompt = template.format(
        context="my context", query_text="user's query", examples=""
    )
    assert (
        prompt
        == """Context:
my context

Examples:


Question:
user's query

Answer:
"""
    )


def test_graphrag_happy_path(retriever_mock: MagicMock, llm: MagicMock) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    retriever_mock.search.return_value = RetrieverResult(
        items=[
            RetrieverResultItem(content="item content 1"),
            RetrieverResultItem(content="item content 2"),
        ]
    )
    llm.invoke.return_value = LLMResponse(content="llm generated text")

    res = rag.search("question", retriever_config={"top_k": 111})

    retriever_mock.search.assert_called_once_with(query_text="question", top_k=111)
    llm.invoke.assert_called_once_with(
        input="""Context:
item content 1
item content 2

Examples:


Question:
question

Answer:
""",
        message_history=None,
        system_instruction="Answer the user question using the provided context.",
    )

    assert isinstance(res, RagResultModel)
    assert res.answer == "llm generated text"
    assert res.retriever_result is None


def test_graphrag_happy_path_with_message_history(
    retriever_mock: MagicMock, llm: MagicMock
) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    retriever_mock.search.return_value = RetrieverResult(
        items=[
            RetrieverResultItem(content="item content 1"),
            RetrieverResultItem(content="item content 2"),
        ]
    )
    llm.invoke.side_effect = [
        LLMResponse(content="llm generated summary"),
        LLMResponse(content="llm generated text"),
    ]
    message_history = [
        {"role": "user", "content": "initial question"},
        {"role": "assistant", "content": "answer to initial question"},
    ]
    res = rag.search("question", message_history)  # type: ignore

    expected_retriever_query_text = """
Message Summary:
llm generated summary

Current Query:
question
"""

    first_invocation_input = """
Summarize the message history:

user: initial question
assistant: answer to initial question
"""
    first_invocation_system_instruction = "You are a summarization assistant. Summarize the given text in no more than 300 words."
    second_invocation = """Context:
item content 1
item content 2

Examples:


Question:
question

Answer:
"""

    retriever_mock.search.assert_called_once_with(
        query_text=expected_retriever_query_text
    )
    assert llm.invoke.call_count == 2
    llm.invoke.assert_has_calls(
        [
            call(
                input=first_invocation_input,
                system_instruction=first_invocation_system_instruction,
            ),
            call(
                input=second_invocation,
                message_history=message_history,
                system_instruction="Answer the user question using the provided context.",
            ),
        ]
    )

    assert isinstance(res, RagResultModel)
    assert res.answer == "llm generated text"
    assert res.retriever_result is None


def test_graphrag_happy_path_with_in_memory_message_history(
    retriever_mock: MagicMock, llm: MagicMock
) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    retriever_mock.search.return_value = RetrieverResult(
        items=[
            RetrieverResultItem(content="item content 1"),
            RetrieverResultItem(content="item content 2"),
        ]
    )
    llm.invoke.side_effect = [
        LLMResponse(content="llm generated summary"),
        LLMResponse(content="llm generated text"),
    ]
    message_history = InMemoryMessageHistory(
        messages=[
            LLMMessage(role="user", content="initial question"),
            LLMMessage(role="assistant", content="answer to initial question"),
        ]
    )
    res = rag.search("question", message_history)

    expected_retriever_query_text = """
Message Summary:
llm generated summary

Current Query:
question
"""

    first_invocation_input = """
Summarize the message history:

user: initial question
assistant: answer to initial question
"""
    first_invocation_system_instruction = "You are a summarization assistant. Summarize the given text in no more than 300 words."
    second_invocation = """Context:
item content 1
item content 2

Examples:


Question:
question

Answer:
"""

    retriever_mock.search.assert_called_once_with(
        query_text=expected_retriever_query_text
    )
    assert llm.invoke.call_count == 2
    llm.invoke.assert_has_calls(
        [
            call(
                input=first_invocation_input,
                system_instruction=first_invocation_system_instruction,
            ),
            call(
                input=second_invocation,
                message_history=message_history.messages,
                system_instruction="Answer the user question using the provided context.",
            ),
        ]
    )

    assert isinstance(res, RagResultModel)
    assert res.answer == "llm generated text"
    assert res.retriever_result is None


def test_graphrag_happy_path_custom_system_instruction(
    retriever_mock: MagicMock, llm: MagicMock
) -> None:
    prompt_template = RagTemplate(system_instructions="Custom instruction")
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
        prompt_template=prompt_template,
    )
    retriever_mock.search.return_value = RetrieverResult(items=[])
    llm.invoke.side_effect = [
        LLMResponse(content="llm generated text"),
    ]
    res = rag.search("question")

    assert llm.invoke.call_count == 1
    llm.invoke.assert_has_calls(
        [
            call(
                input=mock.ANY,
                message_history=None,
                system_instruction="Custom instruction",
            ),
        ]
    )

    assert res.answer == "llm generated text"


def test_graphrag_happy_path_response_fallback(
    retriever_mock: MagicMock, llm: MagicMock
) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    retriever_mock.search.return_value = RetrieverResult(items=[])
    res = rag.search(
        "question",
        response_fallback="I can't answer this question without context",
    )

    assert llm.invoke.call_count == 0
    assert res.answer == "I can't answer this question without context"


def test_graphrag_initialization_error(llm: MagicMock) -> None:
    with pytest.raises(RagInitializationError) as excinfo:
        GraphRAG(
            retriever="not a retriever object",  # type: ignore
            llm=llm,
        )
    assert "retriever" in str(excinfo)


def test_graphrag_search_error(retriever_mock: MagicMock, llm: MagicMock) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    with pytest.raises(SearchValidationError) as excinfo:
        rag.search(10)  # type: ignore
    assert "Input should be a valid string" in str(excinfo)


def test_chat_summary_template(retriever_mock: MagicMock, llm: MagicMock) -> None:
    message_history = [
        {"role": "user", "content": "initial question"},
        {"role": "assistant", "content": "answer to initial question"},
        {"role": "user", "content": "second question"},
        {"role": "assistant", "content": "answer to second question"},
    ]
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    prompt = rag._chat_summary_prompt(message_history=message_history)  # type: ignore
    assert (
        prompt
        == """
Summarize the message history:

user: initial question
assistant: answer to initial question
user: second question
assistant: answer to second question
"""
    )


def test_conversation_template(retriever_mock: MagicMock, llm: MagicMock) -> None:
    rag = GraphRAG(
        retriever=retriever_mock,
        llm=llm,
    )
    prompt = rag.conversation_prompt(
        summary="llm generated chat summary", current_query="latest question"
    )
    assert (
        prompt
        == """
Message Summary:
llm generated chat summary

Current Query:
latest question
"""
    )