# Copyright 2025 Google LLC
#
# 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
#
#     http://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.
#
# SPDX-License-Identifier: Apache-2.0

"""Tests for Anthropic models."""

from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest

from genkit import (
    Media,
    MediaPart,
    Message,
    Metadata,
    ModelConfig,
    ModelRequest,
    ModelResponseChunk,
    Part,
    Role,
    TextPart,
    ToolDefinition,
    ToolRequestPart,
)
from genkit.plugins.anthropic.models import AnthropicModel
from genkit.plugins.anthropic.utils import maybe_strip_fences, strip_markdown_fences


def _create_sample_request() -> ModelRequest:
    """Create a sample generation request for testing."""
    return ModelRequest(
        messages=[
            Message(
                role=Role.USER,
                content=[Part(root=TextPart(text='Hello, how are you?'))],
            )
        ],
        config=ModelConfig(),
        tools=[
            ToolDefinition(
                name='get_weather',
                description='Get weather for a location',
                input_schema={
                    'type': 'object',
                    'properties': {'location': {'type': 'string', 'description': 'Location name'}},
                    'required': ['location'],
                },
            )
        ],
    )


@pytest.mark.asyncio
async def test_generate_basic() -> None:
    """Test basic generation."""
    sample_request = _create_sample_request()

    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_response.content = [MagicMock(type='text', text="Hello! I'm doing well.")]
    mock_response.usage = MagicMock(input_tokens=10, output_tokens=15)
    mock_response.stop_reason = 'end_turn'

    mock_client.messages.create = AsyncMock(return_value=mock_response)

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)
    response = await model.generate(sample_request)

    assert response.message is not None
    assert response.message.content is not None
    assert len(response.message.content) == 1
    part = response.message.content[0]
    actual_part = part.root if isinstance(part, Part) else part
    assert isinstance(actual_part, TextPart)
    assert actual_part.text == "Hello! I'm doing well."
    assert response.usage is not None
    assert response.usage.input_tokens == 10
    assert response.usage.output_tokens == 15
    assert response.finish_reason == 'stop'


@pytest.mark.asyncio
async def test_generate_with_tools() -> None:
    """Test generation with tool calls."""
    sample_request = _create_sample_request()

    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_block = MagicMock()
    mock_block.type = 'tool_use'
    mock_block.id = 'tool_123'
    mock_block.name = 'get_weather'
    mock_block.input = {'location': 'Paris'}
    mock_response.content = [mock_block]
    mock_response.usage = MagicMock(input_tokens=20, output_tokens=10)
    mock_response.stop_reason = 'tool_use'

    mock_client.messages.create = AsyncMock(return_value=mock_response)

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)
    response = await model.generate(sample_request)

    assert response.message is not None
    assert response.message.content is not None
    assert len(response.message.content) == 1
    part = response.message.content[0]
    actual_part = part.root if isinstance(part, Part) else part
    assert isinstance(actual_part, ToolRequestPart)
    assert actual_part.tool_request is not None
    assert actual_part.tool_request.name == 'get_weather'
    assert actual_part.tool_request.ref == 'tool_123'
    assert actual_part.tool_request.input == {'location': 'Paris'}


@pytest.mark.asyncio
async def test_generate_with_config() -> None:
    """Test generation with custom config."""
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_response.content = [MagicMock(type='text', text='Response')]
    mock_response.usage = MagicMock(input_tokens=5, output_tokens=5)
    mock_response.stop_reason = 'end_turn'

    mock_client.messages.create = AsyncMock(return_value=mock_response)

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])],
        config=ModelConfig(
            temperature=0.7,
            max_output_tokens=100,
            top_p=0.9,
        ),
    )

    await model.generate(request)

    call_args = mock_client.messages.create.call_args
    assert call_args.kwargs['temperature'] == 0.7
    assert call_args.kwargs['max_tokens'] == 100
    assert call_args.kwargs['top_p'] == 0.9


def test_extract_system() -> None:
    """Test system prompt extraction."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='You are helpful.'))]),
        Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]),
    ]

    system = model._extract_system(messages)
    assert system == 'You are helpful.'


def test_to_anthropic_messages() -> None:
    """Test message conversion."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]),
        Message(role=Role.MODEL, content=[Part(root=TextPart(text='Hi there'))]),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)

    assert len(anthropic_messages) == 2
    assert anthropic_messages[0]['role'] == 'user'
    assert anthropic_messages[0]['content'][0]['text'] == 'Hello'
    assert anthropic_messages[1]['role'] == 'assistant'
    assert anthropic_messages[1]['content'][0]['text'] == 'Hi there'


class MockStreamManager:
    """Mock stream manager for testing Anthropic streaming."""

    def __init__(self, chunks: list[Any], final_content: list[Any] | None = None) -> None:
        """Initialize the MockStreamManager."""
        self.chunks = chunks
        self.final_message = MagicMock()
        self.final_message.content = final_content if final_content else []
        self.final_message.usage = MagicMock(input_tokens=10, output_tokens=20)
        self.final_message.stop_reason = 'end_turn'

    async def __aenter__(self) -> 'MockStreamManager':
        """Enter the async context manager."""
        return self

    async def __aexit__(self, *args: object) -> None:
        """Exit the async context manager."""
        pass

    def __aiter__(self) -> 'MockStreamManager':
        """Return the async iterator."""
        return self

    async def __anext__(self) -> object:
        """Return the next chunk from the stream."""
        if not self.chunks:
            raise StopAsyncIteration
        return self.chunks.pop(0)

    async def get_final_message(self) -> object:
        """Get the final message from the stream."""
        return self.final_message


@pytest.mark.asyncio
async def test_streaming_generation() -> None:
    """Test streaming generation."""
    sample_request = _create_sample_request()

    mock_client = MagicMock()

    chunks = [
        MagicMock(type='content_block_delta', delta=MagicMock(type='text_delta', text='Hello')),
        MagicMock(type='content_block_delta', delta=MagicMock(type='text_delta', text=' world')),
        MagicMock(type='content_block_delta', delta=MagicMock(type='text_delta', text='!')),
    ]

    final_content = [MagicMock(type='text', text='Hello world!')]
    mock_stream = MockStreamManager(chunks, final_content=final_content)
    mock_client.messages.stream.return_value = mock_stream

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    ctx = MagicMock()
    ctx.is_streaming = True
    collected_chunks: list[ModelResponseChunk] = []

    def send_chunk(chunk: ModelResponseChunk) -> None:
        collected_chunks.append(chunk)

    ctx.send_chunk = send_chunk

    response = await model.generate(sample_request, ctx)

    assert len(collected_chunks) == 3
    chunk0_part = collected_chunks[0].content[0]
    chunk0_actual = chunk0_part.root if isinstance(chunk0_part, Part) else chunk0_part
    assert chunk0_actual.text == 'Hello'

    chunk1_part = collected_chunks[1].content[0]
    chunk1_actual = chunk1_part.root if isinstance(chunk1_part, Part) else chunk1_part
    assert chunk1_actual.text == ' world'

    chunk2_part = collected_chunks[2].content[0]
    chunk2_actual = chunk2_part.root if isinstance(chunk2_part, Part) else chunk2_part
    assert chunk2_actual.text == '!'

    assert response.usage is not None
    assert response.usage.input_tokens == 10
    assert response.usage.output_tokens == 20

    # Verify final response content is populated
    assert response.message is not None
    assert len(response.message.content) == 1
    final_part = response.message.content[0]
    assert isinstance(final_part, Part)
    assert isinstance(final_part.root, TextPart)
    assert final_part.root.text == 'Hello world!'


@pytest.mark.asyncio
async def test_streaming_tool_request() -> None:
    """Test streaming generation with tool use blocks."""
    sample_request = _create_sample_request()

    mock_client = MagicMock()

    # Simulate: text chunk, then tool_use block (start + json deltas + stop).
    tool_block = MagicMock(type='tool_use', id='tool_abc')
    tool_block.name = 'get_weather'
    chunks = [
        MagicMock(type='content_block_delta', delta=MagicMock(type='text_delta', text='Let me check.')),
        MagicMock(type='content_block_start', index=1, content_block=tool_block),
        MagicMock(
            type='content_block_delta',
            index=1,
            delta=MagicMock(type='input_json_delta', partial_json='{"location"'),
        ),
        MagicMock(
            type='content_block_delta',
            index=1,
            delta=MagicMock(type='input_json_delta', partial_json=': "Paris"}'),
        ),
        MagicMock(type='content_block_stop', index=1),
    ]

    final_tool = MagicMock(type='tool_use', id='tool_abc', input={'location': 'Paris'})
    final_tool.name = 'get_weather'
    final_content = [
        MagicMock(type='text', text='Let me check.'),
        final_tool,
    ]
    mock_stream = MockStreamManager(chunks, final_content=final_content)
    mock_client.messages.stream.return_value = mock_stream

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    ctx = MagicMock()
    ctx.is_streaming = True
    collected_chunks: list[ModelResponseChunk] = []
    ctx.send_chunk = lambda chunk: collected_chunks.append(chunk)

    response = await model.generate(sample_request, ctx)

    # Should have 2 chunks: one text, one tool request.
    assert len(collected_chunks) == 2

    text_part = collected_chunks[0].content[0].root
    assert isinstance(text_part, TextPart)
    assert text_part.text == 'Let me check.'

    tool_part = collected_chunks[1].content[0].root
    assert isinstance(tool_part, ToolRequestPart)
    assert tool_part.tool_request.name == 'get_weather'
    assert tool_part.tool_request.ref == 'tool_abc'
    assert tool_part.tool_request.input == {'location': 'Paris'}

    # Final response should also contain the tool request.
    assert response.message is not None
    assert len(response.message.content) == 2


class TestStripMarkdownFences:
    """Tests for strip_markdown_fences."""

    def test_strips_json_fences(self) -> None:
        """Strips ```json ... ``` fences."""
        text = '```json\n{"name": "John", "age": 30}\n```'
        assert strip_markdown_fences(text) == '{"name": "John", "age": 30}'

    def test_strips_plain_fences(self) -> None:
        """Strips ``` ... ``` fences without language tag."""
        text = '```\n{"name": "John"}\n```'
        assert strip_markdown_fences(text) == '{"name": "John"}'

    def test_strips_fences_with_surrounding_whitespace(self) -> None:
        """Strips fences even with leading/trailing whitespace."""
        text = '  \n```json\n{"a": 1}\n```\n  '
        assert strip_markdown_fences(text) == '{"a": 1}'

    def test_preserves_plain_json(self) -> None:
        """Does not alter valid JSON without fences."""
        text = '{"name": "John", "age": 30}'
        assert strip_markdown_fences(text) == text

    def test_preserves_non_json_text(self) -> None:
        """Does not alter plain text."""
        text = 'Hello, world!'
        assert strip_markdown_fences(text) == text

    def test_strips_multiline_json_in_fences(self) -> None:
        """Strips fences around multiline JSON."""
        text = '```json\n{\n  "name": "John",\n  "age": 30\n}\n```'
        result = strip_markdown_fences(text)
        assert result == '{\n  "name": "John",\n  "age": 30\n}'


class TestMaybeStripFences:
    """Tests for maybe_strip_fences."""

    def test_strips_fences_for_json_output(self) -> None:
        """Strips markdown fences when JSON output is requested."""
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='json',
            output_schema={'type': 'object'},
        )
        parts = [Part(root=TextPart(text='```json\n{"a": 1}\n```'))]
        result = maybe_strip_fences(request, parts)
        assert result[0].root.text == '{"a": 1}'

    def test_no_op_for_text_output(self) -> None:
        """Does not modify responses when output format is not json."""
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='text',
        )
        fenced = '```json\n{"a": 1}\n```'
        parts = [Part(root=TextPart(text=fenced))]
        result = maybe_strip_fences(request, parts)
        assert result[0].root.text == fenced

    def test_no_op_for_no_output(self) -> None:
        """Does not modify responses when no output config is set."""
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
        )
        fenced = '```json\n{"a": 1}\n```'
        parts = [Part(root=TextPart(text=fenced))]
        result = maybe_strip_fences(request, parts)
        assert result[0].root.text == fenced

    def test_no_op_when_no_fences(self) -> None:
        """Does not modify clean JSON responses."""
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='json',
            output_schema={'type': 'object'},
        )
        text = '{"name": "John"}'
        parts = [Part(root=TextPart(text=text))]
        result = maybe_strip_fences(request, parts)
        assert result is parts


def test_cache_control_on_text_block() -> None:
    """Test that cache_control metadata is forwarded to Anthropic blocks."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(
            role=Role.USER,
            content=[
                Part(root=TextPart(text='Cached context', metadata=Metadata({'cache_control': {'type': 'ephemeral'}}))),
                Part(root=TextPart(text='Question about the context')),
            ],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)

    assert len(anthropic_messages) == 1
    blocks = anthropic_messages[0]['content']
    assert len(blocks) == 2

    # First block should have cache_control.
    assert blocks[0]['type'] == 'text'
    assert blocks[0]['text'] == 'Cached context'
    assert blocks[0]['cache_control'] == {'type': 'ephemeral'}

    # Second block should not have cache_control.
    assert blocks[1]['type'] == 'text'
    assert 'cache_control' not in blocks[1]


def test_cache_control_not_applied_without_metadata() -> None:
    """Test that no cache_control is applied when metadata is absent."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(
            role=Role.USER,
            content=[Part(root=TextPart(text='No cache'))],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)
    blocks = anthropic_messages[0]['content']
    assert 'cache_control' not in blocks[0]


@pytest.mark.asyncio
async def test_cache_token_tracking_in_usage() -> None:
    """Test that cache creation/read tokens are included in usage."""
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_response.content = [MagicMock(type='text', text='Cached response')]
    mock_response.usage = MagicMock(
        input_tokens=100,
        output_tokens=50,
        cache_creation_input_tokens=80,
        cache_read_input_tokens=20,
    )
    mock_response.stop_reason = 'end_turn'

    mock_client.messages.create = AsyncMock(return_value=mock_response)

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)
    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])],
    )

    response = await model.generate(request)

    assert response.usage is not None
    assert response.usage.input_tokens == 100
    assert response.usage.output_tokens == 50
    assert response.usage.custom is not None
    assert response.usage.custom['cache_creation_input_tokens'] == 80
    assert response.usage.custom['cache_read_input_tokens'] == 20


@pytest.mark.asyncio
async def test_no_cache_tokens_when_caching_not_used() -> None:
    """Test that custom is None when no cache tokens are present."""
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_response.content = [MagicMock(type='text', text='Response')]
    mock_response.usage = MagicMock(input_tokens=10, output_tokens=5)
    mock_response.stop_reason = 'end_turn'
    # Simulate no cache attributes.
    del mock_response.usage.cache_creation_input_tokens
    del mock_response.usage.cache_read_input_tokens

    mock_client.messages.create = AsyncMock(return_value=mock_response)

    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)
    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])],
    )

    response = await model.generate(request)
    assert response.usage is not None
    assert response.usage.custom is None


def test_pdf_base64_becomes_document_block() -> None:
    """Test that a base64 PDF MediaPart converts to Anthropic document block."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    pdf_data = 'data:application/pdf;base64,JVBERi0xLjQ='
    messages = [
        Message(
            role=Role.USER,
            content=[
                Part(root=MediaPart(media=Media(url=pdf_data, content_type='application/pdf'))),
                Part(root=TextPart(text='Summarize this PDF')),
            ],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)
    blocks = anthropic_messages[0]['content']

    assert blocks[0]['type'] == 'document'
    assert blocks[0]['source']['type'] == 'base64'
    assert blocks[0]['source']['media_type'] == 'application/pdf'
    assert blocks[0]['source']['data'] == 'JVBERi0xLjQ='

    assert blocks[1]['type'] == 'text'
    assert blocks[1]['text'] == 'Summarize this PDF'


def test_pdf_url_becomes_document_block() -> None:
    """Test that a URL-based PDF converts to Anthropic document block."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(
            role=Role.USER,
            content=[
                Part(
                    root=MediaPart(
                        media=Media(
                            url='https://example.com/doc.pdf',
                            content_type='application/pdf',
                        )
                    )
                ),
            ],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)
    blocks = anthropic_messages[0]['content']

    assert blocks[0]['type'] == 'document'
    assert blocks[0]['source']['type'] == 'url'
    assert blocks[0]['source']['url'] == 'https://example.com/doc.pdf'


def test_image_still_works() -> None:
    """Test that non-document images still produce image blocks."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    messages = [
        Message(
            role=Role.USER,
            content=[
                Part(root=MediaPart(media=Media(url='https://example.com/cat.jpg', content_type='image/jpeg'))),
            ],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)
    blocks = anthropic_messages[0]['content']

    assert blocks[0]['type'] == 'image'
    assert blocks[0]['source']['type'] == 'url'


def test_pdf_with_cache_control() -> None:
    """Test that cache_control can be applied to document blocks."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    pdf_data = 'data:application/pdf;base64,JVBERi0xLjQ='
    messages = [
        Message(
            role=Role.USER,
            content=[
                Part(
                    root=MediaPart(
                        media=Media(url=pdf_data, content_type='application/pdf'),
                        metadata=Metadata({'cache_control': {'type': 'ephemeral'}}),
                    )
                ),
            ],
        ),
    ]

    anthropic_messages = model._to_anthropic_messages(messages)
    blocks = anthropic_messages[0]['content']

    assert blocks[0]['type'] == 'document'
    assert blocks[0]['cache_control'] == {'type': 'ephemeral'}


def test_structured_output_uses_native_output_config() -> None:
    """Test that JSON schema uses native output_config when model supports it."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-opus-4-6', client=mock_client)

    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a cat'))])],
        output_format='json',
        output_schema={'type': 'object', 'properties': {'name': {'type': 'string'}}},
    )

    params = model._build_params(request)

    assert 'output_config' in params
    assert params['output_config']['format']['type'] == 'json_schema'
    assert params['output_config']['format']['schema']['additionalProperties'] is False


def test_structured_output_falls_back_to_system_prompt() -> None:
    """Test that JSON without schema falls back to system prompt instruction."""
    mock_client = MagicMock()
    model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client)

    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate JSON'))])],
        output_format='json',
    )

    params = model._build_params(request)

    assert 'output_config' not in params
    assert 'system' in params
    assert 'Output valid JSON' in params['system']


def test_structured_output_falls_back_for_unsupported_models() -> None:
    """Test that JSON with schema falls back to system prompt for unsupported models."""
    mock_client = MagicMock()
    # Claude 3.5 Haiku is marked as not supporting JSON natively in model_info.py
    model = AnthropicModel(model_name='claude-3-5-haiku', client=mock_client)

    request = ModelRequest(
        messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a cat'))])],
        output_format='json',
        output_schema={'type': 'object', 'properties': {'name': {'type': 'string'}}},
    )

    params = model._build_params(request)

    assert 'output_config' not in params
    assert 'system' in params
    assert 'Output valid JSON' in params['system']
    assert 'Follow this JSON schema' in params['system']
    assert '"name"' in params['system']
