# 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 OpenAI compatible model implementation."""

from unittest.mock import AsyncMock, MagicMock, PropertyMock

import pytest

from genkit import (
    Message,
    ModelConfig,
    ModelRequest,
    ModelResponse,
    ModelResponseChunk,
    Part,
    Role,
    TextPart,
)
from genkit.plugin_api import ActionRunContext
from genkit.plugins.compat_oai.models import OpenAIModel
from genkit.plugins.compat_oai.models.utils import strip_markdown_fences
from genkit.plugins.compat_oai.typing import OpenAIConfig


def test_get_messages(sample_request: ModelRequest) -> None:
    """Test _get_messages method.

    Ensures the method correctly converts ModelRequest messages into OpenAI-compatible ChatMessage format.
    """
    model = OpenAIModel(model='gpt-4', client=MagicMock())
    messages = model._get_messages(sample_request.messages)

    assert len(messages) == 2
    assert messages[0]['role'] == 'system'
    assert messages[0]['content'] == 'You are an assistant'
    assert messages[1]['role'] == 'user'
    assert messages[1]['content'] == 'Hello, world!'


@pytest.mark.asyncio
async def test_get_openai_config(sample_request: ModelRequest) -> None:
    """Test _get_openai_request_config method.

    Ensures the method correctly constructs the OpenAI API configuration dictionary.
    """
    model = OpenAIModel(model='gpt-4', client=MagicMock())
    openai_config = await model._get_openai_request_config(sample_request)

    assert isinstance(openai_config, dict)
    assert openai_config['model'] == 'gpt-4'
    assert 'messages' in openai_config
    assert isinstance(openai_config['messages'], list)


@pytest.mark.asyncio
async def test__generate(sample_request: ModelRequest) -> None:
    """Test generate method calls OpenAI API and returns ModelResponse."""
    mock_message = MagicMock()
    mock_message.content = 'Hello, user!'
    mock_message.role = 'model'
    mock_message.tool_calls = None
    mock_message.reasoning_content = None

    mock_response = MagicMock()
    mock_response.choices = [MagicMock(message=mock_message)]

    mock_client = MagicMock()
    mock_client.chat.completions.create = AsyncMock(return_value=mock_response)

    model = OpenAIModel(model='gpt-4', client=mock_client)
    response = await model._generate(sample_request)

    mock_client.chat.completions.create.assert_called_once()
    assert isinstance(response, ModelResponse)
    assert response.message is not None
    assert response.message.role == Role.MODEL
    assert response.message.content[0].root.text == 'Hello, user!'


@pytest.mark.asyncio
async def test__generate_stream(sample_request: ModelRequest) -> None:
    """Test generate_stream method ensures it processes streamed responses correctly."""
    mock_client = MagicMock()

    class MockStream:
        def __init__(self, data: list[str]) -> None:
            self._data = data
            self._current = 0

        def __aiter__(self) -> 'MockStream':
            return self

        async def __anext__(self) -> object:
            if self._current >= len(self._data):
                raise StopAsyncIteration

            content = self._data[self._current]
            self._current += 1

            delta_mock = MagicMock()
            delta_mock.content = content
            delta_mock.role = None
            delta_mock.tool_calls = None
            delta_mock.reasoning_content = None

            choice_mock = MagicMock()
            choice_mock.delta = delta_mock

            return MagicMock(choices=[choice_mock])

    mock_client.chat.completions.create = AsyncMock(return_value=MockStream(['Hello', ', world!']))

    model = OpenAIModel(model='gpt-4', client=mock_client)
    collected_chunks = []

    def callback(chunk: ModelResponseChunk) -> None:
        collected_chunks.append(chunk.content[0].root.text)

    await model._generate_stream(sample_request, callback)

    assert collected_chunks == ['Hello', ', world!']


@pytest.mark.parametrize(
    'stream',
    [
        True,
        False,
    ],
)
@pytest.mark.asyncio
async def test_generate(stream: bool, sample_request: ModelRequest) -> None:
    """Tests for generate."""
    ctx_mock = MagicMock(spec=ActionRunContext)
    type(ctx_mock).is_streaming = PropertyMock(return_value=stream)

    mock_response = ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='mocked'))]))

    model = OpenAIModel(model='gpt-4', client=MagicMock())
    model._generate_stream = AsyncMock(return_value=mock_response)
    model._generate = AsyncMock(return_value=mock_response)
    model.normalize_config = MagicMock(return_value={})
    response = await model.generate(sample_request, ctx_mock)

    assert response == mock_response
    if stream:
        model._generate_stream.assert_called_once()
    else:
        model._generate.assert_called_once()


@pytest.mark.parametrize(
    'config, expected',
    [
        (OpenAIConfig(model='test'), OpenAIConfig(model='test')),
        ({'model': 'test'}, OpenAIConfig(model='test')),
        (
            ModelConfig(temperature=0.7),
            OpenAIConfig(temperature=0.7),
        ),
        (
            None,
            Exception(),
        ),
    ],
)
def test_normalize_config(config: object, expected: object) -> None:
    """Tests for _normalize_config."""
    if isinstance(expected, Exception):
        with pytest.raises(ValueError, match=r'Expected request.config to be a dict or OpenAIConfig, got .*'):
            OpenAIModel.normalize_config(config)
    else:
        response = OpenAIModel.normalize_config(config)
        assert response == expected


_SAMPLE_SCHEMA: dict[str, object] = {
    'type': 'object',
    'title': 'RpgCharacter',
    'properties': {
        'name': {'type': 'string'},
        'level': {'type': 'integer'},
    },
    'required': ['name', 'level'],
}


class TestNeedsSchemaInPrompt:
    """Tests for _needs_schema_in_prompt."""

    def test_true_for_deepseek_with_json_and_schema(self) -> None:
        """Returns True for DeepSeek model with json format and schema."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(messages=[], output_format='json', output_schema=_SAMPLE_SCHEMA)
        assert model._needs_schema_in_prompt(request) is True

    def test_false_for_gpt_with_json_and_schema(self) -> None:
        """Returns False for GPT models even with json format and schema."""
        model = OpenAIModel(model='gpt-4o', client=MagicMock())
        request = ModelRequest(messages=[], output_format='json', output_schema=_SAMPLE_SCHEMA)
        assert model._needs_schema_in_prompt(request) is False

    def test_false_for_deepseek_without_schema(self) -> None:
        """Returns False for DeepSeek when no schema is provided."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(messages=[], output_format='json')
        assert model._needs_schema_in_prompt(request) is False

    def test_false_for_deepseek_with_text_format(self) -> None:
        """Returns False for DeepSeek when format is text."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(messages=[], output_format='text')
        assert model._needs_schema_in_prompt(request) is False

    def test_false_for_no_format(self) -> None:
        """Returns False when output has no format set."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(messages=[])
        assert model._needs_schema_in_prompt(request) is False


class TestBuildSchemaInstruction:
    """Tests for _build_schema_instruction."""

    def test_returns_system_message(self) -> None:
        """Returns a dict with role 'system'."""
        result = OpenAIModel._build_schema_instruction(_SAMPLE_SCHEMA)
        assert result['role'] == 'system'

    def test_content_contains_schema(self) -> None:
        """Content includes the schema's field names and title."""
        result = OpenAIModel._build_schema_instruction(_SAMPLE_SCHEMA)
        assert '"RpgCharacter"' in result['content']
        assert '"name"' in result['content']
        assert '"level"' in result['content']

    def test_content_contains_instructions(self) -> None:
        """Content includes directive keywords."""
        result = OpenAIModel._build_schema_instruction(_SAMPLE_SCHEMA)
        assert 'EXACTLY' in result['content']
        assert 'JSON schema' in result['content']


class TestSchemaInjectionInConfig:
    """Tests for schema injection in _get_openai_request_config."""

    @pytest.mark.asyncio
    async def test_deepseek_injects_schema_message(self) -> None:
        """DeepSeek request prepends a schema system message."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[
                Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a character'))]),
            ],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        config = await model._get_openai_request_config(request)

        messages = config['messages']
        # Schema instruction is prepended as the first message.
        assert messages[0]['role'] == 'system'
        assert 'RpgCharacter' in messages[0]['content']
        # Original user message follows.
        assert messages[1]['role'] == 'user'
        assert messages[1]['content'] == 'Generate a character'

    @pytest.mark.asyncio
    async def test_gpt_does_not_inject_schema_message(self) -> None:
        """GPT request does not prepend a schema system message."""
        model = OpenAIModel(model='gpt-4o', client=MagicMock())
        request = ModelRequest(
            messages=[
                Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a character'))]),
            ],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        config = await model._get_openai_request_config(request)

        messages = config['messages']
        # No extra system message — only the original user message.
        assert len(messages) == 1
        assert messages[0]['role'] == 'user'

    @pytest.mark.asyncio
    async def test_deepseek_without_schema_no_injection(self) -> None:
        """DeepSeek request without a schema does not inject anything."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[
                Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]),
            ],
            output_format='json',
        )
        config = await model._get_openai_request_config(request)

        messages = config['messages']
        assert len(messages) == 1
        assert messages[0]['role'] == 'user'

    @pytest.mark.asyncio
    async def test_deepseek_preserves_existing_system_message(self) -> None:
        """Schema injection does not clobber an existing system message."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[
                Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='You are helpful'))]),
                Message(role=Role.USER, content=[Part(root=TextPart(text='Generate'))]),
            ],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        config = await model._get_openai_request_config(request)

        messages = config['messages']
        # Schema instruction prepended, then original system, then user.
        assert len(messages) == 3
        assert messages[0]['role'] == 'system'
        assert 'RpgCharacter' in messages[0]['content']
        assert messages[1]['role'] == 'system'
        assert messages[1]['content'] == 'You are helpful'
        assert messages[2]['role'] == 'user'


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 TestCleanJsonResponse:
    """Tests for _clean_json_response."""

    def test_cleans_deepseek_json_response(self) -> None:
        """Strips markdown fences from DeepSeek JSON response."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        response = ModelResponse(
            request=request,
            message=Message(
                role=Role.MODEL,
                content=[Part(root=TextPart(text='```json\n{"name": "John", "level": 5}\n```'))],
            ),
        )
        cleaned = model._clean_json_response(response, request)
        assert cleaned.message is not None
        assert cleaned.message.content[0].root.text == '{"name": "John", "level": 5}'

    def test_no_op_for_gpt_model(self) -> None:
        """Does not modify responses from non-DeepSeek models."""
        model = OpenAIModel(model='gpt-4o', client=MagicMock())
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        fenced_text = '```json\n{"name": "John", "level": 5}\n```'
        response = ModelResponse(
            request=request,
            message=Message(
                role=Role.MODEL,
                content=[Part(root=TextPart(text=fenced_text))],
            ),
        )
        result = model._clean_json_response(response, request)
        assert result.message is not None
        assert result.message.content[0].root.text == fenced_text

    def test_no_op_for_text_output(self) -> None:
        """Does not modify responses when output format is not json."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='text',
        )
        text = '```json\n{"a": 1}\n```'
        response = ModelResponse(
            request=request,
            message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]),
        )
        result = model._clean_json_response(response, request)
        assert result.message is not None
        assert result.message.content[0].root.text == text

    def test_no_op_for_no_output(self) -> None:
        """Does not modify responses when no output config is set."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
        )
        text = '```json\n{"a": 1}\n```'
        response = ModelResponse(
            request=request,
            message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]),
        )
        result = model._clean_json_response(response, request)
        assert result.message is not None
        assert result.message.content[0].root.text == text

    def test_no_op_when_no_fences(self) -> None:
        """Does not modify clean JSON responses."""
        model = OpenAIModel(model='deepseek-chat', client=MagicMock())
        request = ModelRequest(
            messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])],
            output_format='json',
            output_schema=_SAMPLE_SCHEMA,
        )
        text = '{"name": "John", "level": 5}'
        response = ModelResponse(
            request=request,
            message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]),
        )
        result = model._clean_json_response(response, request)
        # Should return the exact same object (no copy).
        assert result is response
