#  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 __future__ import annotations
import warnings
import sys
from typing import Any, Generator, List, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch

import anthropic
import pytest
from neo4j_graphrag.exceptions import LLMGenerationError
from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM
from neo4j_graphrag.llm.types import LLMResponse
from neo4j_graphrag.types import LLMMessage
from pydantic import BaseModel, ConfigDict


@pytest.fixture
def mock_anthropic() -> Generator[MagicMock, None, None]:
    mock = MagicMock()
    mock.APIError = anthropic.APIError
    mock.NOT_GIVEN = anthropic.NOT_GIVEN

    with patch.dict(sys.modules, {"anthropic": mock}):
        yield mock


def _as_mock(value: Any) -> MagicMock:
    return cast(MagicMock, value)


def _as_async_mock(value: Any) -> AsyncMock:
    return cast(AsyncMock, value)


@patch("builtins.__import__", side_effect=ImportError)
def test_anthropic_llm_missing_dependency(mock_import: Mock) -> None:
    with pytest.raises(ImportError):
        AnthropicLLM(model_name="claude-3-opus-20240229")


def test_anthropic_invoke_happy_path(mock_anthropic: Mock) -> None:
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="generated text")]
    )
    model_params = {"temperature": 0.3}
    llm = AnthropicLLM("claude-3-opus-20240229", model_params=model_params)
    input_text = "may thy knife chip and shatter"
    response = llm.invoke(input_text)
    assert response.content == "generated text"
    _as_mock(llm.client.messages.create).assert_called_once_with(
        messages=[{"role": "user", "content": input_text}],
        model="claude-3-opus-20240229",
        system=anthropic.NOT_GIVEN,
        **model_params,
    )


def test_anthropic_invoke_with_message_history_happy_path(mock_anthropic: Mock) -> None:
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="generated text")]
    )
    model_params = {"temperature": 0.3}
    llm = AnthropicLLM(
        "claude-3-opus-20240229",
        model_params=model_params,
    )
    message_history = [
        {"role": "user", "content": "When does the sun come up in the summer?"},
        {"role": "assistant", "content": "Usually around 6am."},
    ]
    question = "What about next season?"

    response = llm.invoke(question, message_history)  # type: ignore
    assert response.content == "generated text"
    message_history.append({"role": "user", "content": question})
    _as_mock(llm.client.messages.create).assert_called_once_with(
        messages=message_history,
        model="claude-3-opus-20240229",
        system=anthropic.NOT_GIVEN,
        **model_params,
    )


def test_anthropic_invoke_with_system_instruction(
    mock_anthropic: Mock,
) -> None:
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="generated text")]
    )
    model_params = {"temperature": 0.3}
    system_instruction = "You are a helpful assistant."
    llm = AnthropicLLM(
        "claude-3-opus-20240229",
        model_params=model_params,
    )

    question = "When does it come up in the winter?"
    response = llm.invoke(question, system_instruction=system_instruction)
    assert isinstance(response, LLMResponse)
    assert response.content == "generated text"
    messages: List[LLMMessage] = [{"role": "user", "content": question}]
    _as_mock(llm.client.messages.create).assert_called_with(
        model="claude-3-opus-20240229",
        system=system_instruction,
        messages=messages,
        **model_params,
    )

    assert _as_mock(llm.client.messages.create).call_count == 1


def test_anthropic_invoke_with_message_history_and_system_instruction(
    mock_anthropic: Mock,
) -> None:
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="generated text")]
    )
    model_params = {"temperature": 0.3}
    system_instruction = "You are a helpful assistant."
    llm = AnthropicLLM(
        "claude-3-opus-20240229",
        model_params=model_params,
    )
    message_history = [
        {"role": "user", "content": "When does the sun come up in the summer?"},
        {"role": "assistant", "content": "Usually around 6am."},
    ]

    question = "When does it come up in the winter?"
    response = llm.invoke(question, message_history, system_instruction)  # type: ignore
    assert isinstance(response, LLMResponse)
    assert response.content == "generated text"
    message_history.append({"role": "user", "content": question})
    _as_mock(llm.client.messages.create).assert_called_with(
        model="claude-3-opus-20240229",
        system=system_instruction,
        messages=message_history,
        **model_params,
    )

    assert _as_mock(llm.client.messages.create).call_count == 1


def test_anthropic_invoke_with_message_history_validation_error(
    mock_anthropic: Mock,
) -> None:
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="generated text")]
    )
    model_params = {"temperature": 0.3}
    system_instruction = "You are a helpful assistant."
    llm = AnthropicLLM(
        "claude-3-opus-20240229",
        model_params=model_params,
        system_instruction=system_instruction,
    )
    message_history = [
        {"role": "human", "content": "When does the sun come up in the summer?"},
        {"role": "assistant", "content": "Usually around 6am."},
    ]
    question = "What about next season?"

    with pytest.raises(LLMGenerationError) as exc_info:
        llm.invoke(question, message_history)  # type: ignore
    assert "Input should be 'user', 'assistant' or 'system'" in str(exc_info.value)


@pytest.mark.asyncio
async def test_anthropic_ainvoke_happy_path(mock_anthropic: Mock) -> None:
    mock_response = AsyncMock()
    mock_response.content = [MagicMock(text="Return text")]
    mock_model = mock_anthropic.AsyncAnthropic.return_value
    mock_model.messages.create = AsyncMock(return_value=mock_response)
    model_params = {"temperature": 0.3}
    llm = AnthropicLLM("claude-3-opus-20240229", model_params)
    input_text = "may thy knife chip and shatter"
    response = await llm.ainvoke(input_text)
    assert response.content == "Return text"
    _as_async_mock(llm.async_client.messages.create).assert_awaited_once_with(
        model="claude-3-opus-20240229",
        system=anthropic.NOT_GIVEN,
        messages=[{"role": "user", "content": input_text}],
        **model_params,
    )


# V2 Interface Tests


def test_anthropic_llm_invoke_v2_happy_path(mock_anthropic: Mock) -> None:
    """Test V2 interface invoke method with List[LLMMessage] input."""
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="anthropic v2 response")]
    )
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is machine learning?"},
    ]

    model_params = {"temperature": 0.7}
    llm = AnthropicLLM(model_name="claude-3-opus-20240229", model_params=model_params)
    response = llm.invoke(messages)

    assert isinstance(response, LLMResponse)
    assert response.content == "anthropic v2 response"

    # Verify the correct method was called with system instruction and messages
    _as_mock(llm.client.messages.create).assert_called_once_with(
        model="claude-3-opus-20240229",
        system="You are a helpful assistant.",
        messages=[{"role": "user", "content": "What is machine learning?"}],
        **model_params,
    )


def test_anthropic_llm_invoke_v2_with_conversation_history(
    mock_anthropic: Mock,
) -> None:
    """Test V2 interface invoke method with complex conversation history."""
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="anthropic conversation response")]
    )
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me about Python."},
        {"role": "assistant", "content": "Python is a programming language."},
        {"role": "user", "content": "What about its history?"},
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
    response = llm.invoke(messages)

    assert isinstance(response, LLMResponse)
    assert response.content == "anthropic conversation response"

    # Verify the correct number of messages were passed (excluding system)
    _as_mock(llm.client.messages.create).assert_called_once()
    call_args = _as_mock(llm.client.messages.create).call_args[1]
    assert call_args["system"] == "You are a helpful assistant."
    assert len(call_args["messages"]) == 3


def test_anthropic_llm_invoke_v2_no_system_message(mock_anthropic: Mock) -> None:
    """Test V2 interface invoke method without system message."""
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[MagicMock(text="anthropic no system response")]
    )
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "user", "content": "What is the capital of France?"},
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
    response = llm.invoke(messages)

    assert isinstance(response, LLMResponse)
    assert response.content == "anthropic no system response"

    # Verify only user message was passed and no system instruction
    _as_mock(llm.client.messages.create).assert_called_once()
    call_args = _as_mock(llm.client.messages.create).call_args[1]
    assert call_args["system"] == anthropic.NOT_GIVEN
    assert len(call_args["messages"]) == 1


@pytest.mark.asyncio
async def test_anthropic_llm_ainvoke_v2_happy_path(mock_anthropic: Mock) -> None:
    """Test V2 interface async invoke method with List[LLMMessage] input."""
    mock_response = AsyncMock()
    mock_response.content = [MagicMock(text="async anthropic v2 response")]
    mock_model = mock_anthropic.AsyncAnthropic.return_value
    mock_model.messages.create = AsyncMock(return_value=mock_response)
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is async programming?"},
    ]

    model_params = {"max_tokens": 100}
    llm = AnthropicLLM(model_name="claude-3-opus-20240229", model_params=model_params)
    response = await llm.ainvoke(messages)

    assert isinstance(response, LLMResponse)
    assert response.content == "async anthropic v2 response"

    # Verify the async client was called correctly
    _as_async_mock(llm.async_client.messages.create).assert_awaited_once_with(
        model="claude-3-opus-20240229",
        system="You are a helpful assistant.",
        messages=[{"role": "user", "content": "What is async programming?"}],
        **model_params,
    )


def test_anthropic_llm_invoke_v2_validation_error(mock_anthropic: Mock) -> None:
    """Test V2 interface invoke method with invalid role."""
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "invalid_role", "content": "This should fail."},  # type: ignore[typeddict-item]
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")

    with pytest.raises(ValueError) as exc_info:
        llm.invoke(messages)
    assert "Unknown role: invalid_role" in str(exc_info.value)


def test_anthropic_llm_invoke_invalid_input_type(
    mock_anthropic: Mock,
) -> None:  # noqa: ARG001
    """Test that invalid input type raises appropriate error."""
    llm = AnthropicLLM(model_name="claude-3-opus-20240229")

    with pytest.raises(ValueError) as exc_info:
        llm.invoke(123)  # type: ignore
    assert "Invalid input type for invoke method" in str(exc_info.value)


@pytest.mark.asyncio
async def test_anthropic_llm_ainvoke_invalid_input_type(
    mock_anthropic: Mock,
) -> None:  # noqa: ARG001
    """Test that invalid input type raises appropriate error for async invoke."""
    llm = AnthropicLLM(model_name="claude-3-opus-20240229")

    with pytest.raises(ValueError) as exc_info:
        await llm.ainvoke(123)  # type: ignore
    assert "Invalid input type for ainvoke method" in str(exc_info.value)


def test_anthropic_llm_get_brand_new_messages_all_roles(mock_anthropic: Mock) -> None:
    """Test get_brand_new_messages method handles all message roles correctly."""

    def create_message_param(**kwargs: str) -> MagicMock:
        return MagicMock(**kwargs)

    mock_anthropic.types.MessageParam = MagicMock(side_effect=create_message_param)

    messages: List[LLMMessage] = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there!"},
        {"role": "user", "content": "How are you?"},
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
    system_instruction, result_messages = llm.get_messages_v2(messages)

    # Verify system instruction is extracted
    assert system_instruction == "You are a helpful assistant."

    result_messages = cast(list[MagicMock], list(result_messages))

    # Verify the correct number of non-system messages are returned
    assert len(result_messages) == 3

    # Verify message content is preserved
    assert result_messages[0].content == "Hello"
    assert result_messages[1].content == "Hi there!"
    assert result_messages[2].content == "How are you?"


def test_anthropic_llm_get_brand_new_messages_unknown_role(
    mock_anthropic: Mock,
) -> None:  # noqa: ARG001
    """Test get_brand_new_messages method raises error for unknown role."""
    messages: List[LLMMessage] = [
        {"role": "unknown_role", "content": "This should fail."},  # type: ignore[typeddict-item]
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")

    with pytest.raises(ValueError) as exc_info:
        llm.get_messages_v2(messages)
    assert "Unknown role: unknown_role" in str(exc_info.value)


def test_anthropic_llm_invoke_v2_empty_response_error(mock_anthropic: Mock) -> None:
    """Test V2 interface invoke method handles empty response."""
    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
        content=[]  # Empty content should trigger error
    )
    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)

    messages: List[LLMMessage] = [
        {"role": "user", "content": "This should return empty response."},
    ]

    llm = AnthropicLLM(model_name="claude-3-opus-20240229")

    with pytest.raises(LLMGenerationError) as exc_info:
        llm.invoke(messages)
    assert "LLM returned empty response" in str(exc_info.value)


def test_anthropic_invoke_v2_with_response_format_raises_error(
    mock_anthropic: Mock,
) -> None:
    """Test V2 interface raises NotImplementedError when response_format is used."""

    class TestModel(BaseModel):
        model_config = ConfigDict(extra="forbid")
        value: str

    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
    llm = AnthropicLLM(api_key="test", model_name="claude-3-opus")

    with pytest.raises(NotImplementedError) as exc_info:
        llm.invoke(messages, response_format=TestModel)

    assert "AnthropicLLM does not currently support structured output" in str(
        exc_info.value
    )


def test_anthropic_llm_close(mock_anthropic: Mock) -> None:
    mock_anthropic.AsyncAnthropic.return_value.close = AsyncMock()

    llm = AnthropicLLM("claude-3-opus-20240229")

    with warnings.catch_warnings():
        warnings.simplefilter("error")
        llm.close()

    mock_anthropic.Anthropic.return_value.close.assert_called_once()
    mock_anthropic.AsyncAnthropic.return_value.close.assert_called_once()


@pytest.mark.asyncio
async def test_anthropic_llm_aclose(mock_anthropic: Mock) -> None:
    mock_anthropic.AsyncAnthropic.return_value.close = AsyncMock()

    llm = AnthropicLLM("claude-3-opus-20240229")

    with warnings.catch_warnings():
        warnings.simplefilter("error")
        await llm.aclose()

    mock_anthropic.Anthropic.return_value.close.assert_called_once()
    mock_anthropic.AsyncAnthropic.return_value.close.assert_called_once()