# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import json
from typing import Any
from unittest.mock import MagicMock

import pytest

from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.tool_parsers.apertus_tool_parser import (
    TOOL_CALLS_PREFIX,
    TOOL_CALLS_SUFFIX,
    ApertusToolParser,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def mock_tokenizer():
    tokenizer = MagicMock()
    tokenizer.encode.return_value = [1, 2, 3]
    # Include the tool call tokens in the vocab for the parser
    tokenizer.get_vocab.return_value = {TOOL_CALLS_PREFIX: 100, TOOL_CALLS_SUFFIX: 101}
    return tokenizer


@pytest.fixture
def parser(mock_tokenizer):
    return ApertusToolParser(mock_tokenizer)


@pytest.fixture
def mock_request():
    request = MagicMock(spec=ChatCompletionRequest)
    request.tools = []
    request.tool_choice = "auto"
    return request


# ---------------------------------------------------------------------------
# Non-streaming extraction tests
# ---------------------------------------------------------------------------


class TestExtractToolCalls:
    def test_no_tool_calls(self, parser, mock_request):
        model_output = "Hello, how can I help you today?"
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is False
        assert result.tool_calls == []
        assert result.content == model_output

    def test_single_tool_call(self, parser, mock_request):
        model_output = (
            '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "get_weather"
        args = json.loads(result.tool_calls[0].function.arguments)
        assert args == {"location": "London"}

    def test_multiple_arguments(self, parser, mock_request):
        model_output = (
            '<|tools_prefix|>[{"get_weather": '
            '{"location": "San Francisco", '
            '"unit": "celsius"}}]<|tools_suffix|>'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "get_weather"
        args = json.loads(result.tool_calls[0].function.arguments)
        assert args == {"location": "San Francisco", "unit": "celsius"}

    def test_text_before_tool_call(self, parser, mock_request):
        model_output = (
            "Let me check the weather for you. "
            '<|tools_prefix|>[{"get_weather": {"location": "Paris"}}]<|tools_suffix|>'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert result.content == "Let me check the weather for you."
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "get_weather"

    def test_multiple_tool_calls(self, parser, mock_request):
        model_output = (
            '<|tools_prefix|>[{"get_weather": '
            '{"location": "London"}}, '
            '{"get_time": {"location": "London"}}]<|tools_suffix|>'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 2
        assert result.tool_calls[0].function.name == "get_weather"
        assert result.tool_calls[1].function.name == "get_time"

    def test_nested_arguments(self, parser, mock_request):
        model_output = (
            '<|tools_prefix|>[{"complex_function": '
            '{"nested": {"inner": "value"}, '
            '"list": ["a", "b"]}}]<|tools_suffix|>'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "complex_function"
        args = json.loads(result.tool_calls[0].function.arguments)
        assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]}

    def test_incomplete_tool_call(self, parser, mock_request):
        model_output = '<|tools_prefix|>[{"get_weather": {"location": "London"}'
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "get_weather"
        args = json.loads(result.tool_calls[0].function.arguments)
        assert args == {"location": "London"}

    def test_missing_tool_suffix(self, parser, mock_request):
        model_output = (
            '<|tools_prefix|>[{"get_weather": '
            '{"location": "San Francisco", "unit": "celsius"}}]'
        )
        result = parser.extract_tool_calls(model_output, mock_request)

        assert result.tools_called is True
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "get_weather"
        args = json.loads(result.tool_calls[0].function.arguments)
        assert args == {"location": "San Francisco", "unit": "celsius"}


# ---------------------------------------------------------------------------
# Streaming extraction tests
# ---------------------------------------------------------------------------


class TestStreamingExtraction:
    def _simulate_streaming(
        self, parser: ApertusToolParser, mock_request: Any, chunks: list[str]
    ) -> list[tuple[Any, str]]:
        results: list[tuple[Any, str]] = []
        previous_text: str = ""
        previous_token_ids: list[int] = []

        for chunk in chunks:
            current_text = previous_text + chunk
            # Simulate a token ID sequence matching the chunk progression
            delta_token_ids: list[int] = [0] * max(1, len(chunk) // 4)
            current_token_ids = previous_token_ids + delta_token_ids

            delta = parser.extract_tool_calls_streaming(
                previous_text=previous_text,
                current_text=current_text,
                delta_text=chunk,
                previous_token_ids=tuple(previous_token_ids),
                current_token_ids=tuple(current_token_ids),
                delta_token_ids=tuple(delta_token_ids),
                request=mock_request,
            )
            results.append((delta, current_text))
            previous_text = current_text
            previous_token_ids = list(current_token_ids)

        return results

    def _collect_tool_calls(self, results) -> dict[int, dict[str, Any]]:
        """Properly tracks and concatenates streamed tool arguments by their Index."""
        tool_calls = {}
        for delta, _ in results:
            if not delta or not getattr(delta, "tool_calls", None):
                continue

            for tc in delta.tool_calls:
                idx = (
                    tc.get("index", 0)
                    if isinstance(tc, dict)
                    else getattr(tc, "index", 0)
                )
                func = (
                    tc.get("function", {})
                    if isinstance(tc, dict)
                    else getattr(tc, "function", None)
                )
                if not func:
                    continue

                name = (
                    func.get("name")
                    if isinstance(func, dict)
                    else getattr(func, "name", None)
                )
                args = (
                    func.get("arguments")
                    if isinstance(func, dict)
                    else getattr(func, "arguments", None)
                )

                if idx not in tool_calls:
                    tool_calls[idx] = {"name": "", "arguments": ""}

                if name:
                    tool_calls[idx]["name"] += name
                if args:
                    tool_calls[idx]["arguments"] += args

        return tool_calls

    def _collect_content(self, results) -> str:
        """Collects generated normal text outside of the tool calls."""
        return "".join(
            delta.content
            for delta, _ in results
            if delta and getattr(delta, "content", None)
        )

    def test_basic_streaming_single_tool(self, parser, mock_request):
        chunks = [
            "<|tools_prefix|>",
            '[{"get_weather": ',
            '{"location": "Paris, ',
            'France"}}]',
            "<|tools_suffix|>",
        ]

        results = self._simulate_streaming(parser, mock_request, chunks)
        tcs = self._collect_tool_calls(results)

        assert len(tcs) == 1
        assert tcs[0]["name"] == "get_weather"
        assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"}

    def test_streaming_missing_tool_suffix(self, parser, mock_request):
        chunks = [
            "<|tools_prefix|>",
            '[{"get_weather": ',
            '{"location": "Paris, ',
            'France"}}]',
        ]

        results = self._simulate_streaming(parser, mock_request, chunks)
        tcs = self._collect_tool_calls(results)

        assert len(tcs) == 1
        assert tcs[0]["name"] == "get_weather"
        assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"}

    def test_streaming_partial_tag_buffering_missing_tool_suffix(
        self, parser, mock_request
    ):
        chunks = ["Content", "<|tools_", "prefix|>", '[{"f": ', '{"a": 1}}]']

        results = self._simulate_streaming(parser, mock_request, chunks)
        content = self._collect_content(results)

        assert "Content" in content
        assert "<|tools_prefix|>" not in content
        assert "<|tools_suffix|>" not in content

        tcs = self._collect_tool_calls(results)

        assert len(tcs) == 1
        assert tcs[0]["name"] == "f"
        assert json.loads(tcs[0]["arguments"]) == {"a": 1}

    def test_streaming_multi_tool(self, parser, mock_request):
        chunks = [
            "<|tools_prefix|>",
            '[{"get_weather": {"location": "Tokyo"}}',
            ', {"get_time": {"location": "Tokyo"}}]',
            "<|tools_suffix|>",
        ]

        results = self._simulate_streaming(parser, mock_request, chunks)
        tcs = self._collect_tool_calls(results)

        assert len(tcs) == 2
        assert tcs[0]["name"] == "get_weather"
        assert json.loads(tcs[0]["arguments"]) == {"location": "Tokyo"}
        assert tcs[1]["name"] == "get_time"
        assert json.loads(tcs[1]["arguments"]) == {"location": "Tokyo"}

    def test_streaming_text_before_tool_call(self, parser, mock_request):
        chunks = [
            "Let me check ",
            "the weather. ",
            "<|tools_prefix|>",
            '[{"get_weather": {"location": "London"}}]',
            "<|tools_suffix|>",
        ]

        results = self._simulate_streaming(parser, mock_request, chunks)
        content = self._collect_content(results)

        assert content.strip() == "Let me check the weather."
        tcs = self._collect_tool_calls(results)

        assert len(tcs) == 1
        assert tcs[0]["name"] == "get_weather"
        assert json.loads(tcs[0]["arguments"]) == {"location": "London"}

    def test_streaming_partial_tag_buffering(self, parser, mock_request):
        chunks = [
            "Content",
            "<|tools_",
            "prefix|>",
            '[{"f": {"a": 1}}]',
            "<|tools_suf",
            "fix|>",
        ]

        results = self._simulate_streaming(parser, mock_request, chunks)
        content = self._collect_content(results)

        assert "Content" in content
        assert "<|tools_prefix|>" not in content
        assert "<|tools_suffix|>" not in content

        tc = self._collect_tool_calls(results)
        assert len(tc) == 1
        assert tc[0]["name"] == "f"
        assert json.loads(tc[0]["arguments"]) == {"a": 1}

    # ---------------------------------------------------------------------------
    # Edge Cases: Multi-Token Prediction (MTP) & vLLM Chunking Anomalies
    # ---------------------------------------------------------------------------

    def test_mtp_streaming_massive_chunk(self, parser, mock_request):
        """Simulates MTP predicting text, tool calls,
        and trailing text all in a single chunk."""
        chunks = [
            "Sure! "
            '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>'
        ]
        results = self._simulate_streaming(parser, mock_request, chunks)

        content = self._collect_content(results)
        assert "Sure! " in content

        tc = self._collect_tool_calls(results)
        assert len(tc) == 1
        assert tc[0]["name"] == "get_weather"
        assert json.loads(tc[0]["arguments"]) == {"location": "London"}

    def test_mtp_streaming_multiple_tools_burst(self, parser, mock_request):
        """Simulates MTP predicting an array of multiple tools in one single chunk."""
        chunks = [
            '<|tools_prefix|>[{"get_weather": '
            '{"location": "London"}}, '
            '{"get_time": {"location": "Paris"}}]<|tools_suffix|>'
        ]
        results = self._simulate_streaming(parser, mock_request, chunks)

        tc = self._collect_tool_calls(results)
        assert len(tc) == 2
        assert tc[0]["name"] == "get_weather"
        assert json.loads(tc[0]["arguments"]) == {"location": "London"}
        assert tc[1]["name"] == "get_time"
        assert json.loads(tc[1]["arguments"]) == {"location": "Paris"}

    def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request):
        """Simulates MTP chunks that jump over entire tools
        (e.g., from middle of tool 1 to middle of tool 3)."""
        chunks = [
            '<|tools_prefix|>[{"t1": {"a": 1}',
            '}, {"t2": {"b": 2}}, {"t3": {"c": 3',
            "}}]<|tools_suffix|>",
        ]
        results = self._simulate_streaming(parser, mock_request, chunks)

        tc = self._collect_tool_calls(results)
        assert len(tc) == 3
        assert tc[0]["name"] == "t1"
        assert json.loads(tc[0]["arguments"]) == {"a": 1}
        assert tc[1]["name"] == "t2"
        assert json.loads(tc[1]["arguments"]) == {"b": 2}
        assert tc[2]["name"] == "t3"
        assert json.loads(tc[2]["arguments"]) == {"c": 3}

    def test_vllm_streaming_character_by_character(self, parser, mock_request):
        """Simulates worst-case vLLM fragmentation where
        chunks arrive character-by-character."""
        text = (
            'Hi <|tools_prefix|>[{"get_weather": '
            '{"location": "London"}}]<|tools_suffix|> '
        )
        chunks = list(text)
        results = self._simulate_streaming(parser, mock_request, chunks)

        content = self._collect_content(results)
        assert "Hi" in content

        tc = self._collect_tool_calls(results)
        assert len(tc) == 1
        assert tc[0]["name"] == "get_weather"
        assert json.loads(tc[0]["arguments"]) == {"location": "London"}

    def test_vllm_streaming_empty_deltas(self, parser, mock_request):
        """Simulates vLLM stream producing empty string chunks
        (e.g., hidden tokens or artifacts)."""
        chunks = [
            "Wait",
            "",
            "<|tools_prefix|>",
            "",
            '[{"get_weather": ',
            "",
            '{"location": "London"}}]',
            "<|tools_suffix|>",
        ]
        results = self._simulate_streaming(parser, mock_request, chunks)

        content = self._collect_content(results)
        assert content == "Wait"

        tc = self._collect_tool_calls(results)
        assert len(tc) == 1
        assert tc[0]["name"] == "get_weather"
        assert json.loads(tc[0]["arguments"]) == {"location": "London"}