# pyright: reportPrivateUsage=false

import json
from typing import Any, cast

import pytest
from httpx import Response
from pydantic import BaseModel

from core.domain.message import Message, MessageContent
from core.domain.models import Model, Provider
from core.domain.tool_call import ToolCallRequest, ToolCallResult
from core.providers._base.abstract_provider import RawCompletion
from core.providers._base.llm_usage import LLMUsage
from core.providers._base.provider_error import ModelDoesNotSupportModeError, ProviderInvalidFileError
from core.providers._base.provider_options import ProviderOptions
from core.providers.openai.openai_domain import (
    ChoiceDelta,
    CompletionRequest,
    JSONResponseFormat,
    JSONSchemaResponseFormat,
    OpenAIError,
    StreamedResponse,
    StreamedToolCall,
    StreamedToolCallFunction,
    TextResponseFormat,
    Usage,
)
from core.providers.openai.openai_provider_base import OpenAIProviderBase, OpenAIProviderBaseConfig
from core.runners.runner_output import ToolCallRequestDelta
from tests import fake_models as test_models
from tests.utils import fixtures_json


class _TestProviderConfig(OpenAIProviderBaseConfig):
    """Test implementation of OpenAIProviderBaseConfig."""

    @property
    def provider(self) -> Provider:
        return Provider.OPEN_AI


class _TestOpenAIProviderBase(OpenAIProviderBase[_TestProviderConfig]):
    """Test implementation of OpenAIProviderBase."""

    @classmethod
    def required_env_vars(cls) -> list[str]:
        return []

    @classmethod
    def name(cls) -> Provider:
        return Provider.OPEN_AI

    @classmethod
    def _default_config(cls, index: int) -> _TestProviderConfig:
        return _TestProviderConfig()

    def _request_url(self, model: Model, stream: bool) -> str:
        return "test"

    async def _request_headers(self, request: dict[str, Any], url: str, model: Model) -> dict[str, str]:
        return {}


@pytest.fixture
def base_provider() -> _TestOpenAIProviderBase:
    return _TestOpenAIProviderBase()


def test_extract_stream_delta_max_tokens_exceeded() -> None:
    """Test that length finish_reason raises MaxTokensExceededError."""
    provider = _TestOpenAIProviderBase()

    event = (
        StreamedResponse(
            choices=[
                ChoiceDelta(
                    index=0,
                    delta=ChoiceDelta.MessageDelta(content="test"),
                    finish_reason="length",
                ),
            ],
        )
        .model_dump_json()
        .encode()
    )

    res = provider._extract_stream_delta(event)
    assert res.finish_reason == "max_context"


def test_extract_stream_delta_with_content() -> None:
    """Test successful extraction of content from stream delta."""
    provider = _TestOpenAIProviderBase()

    event = (
        StreamedResponse(
            choices=[
                ChoiceDelta(
                    index=0,
                    delta=ChoiceDelta.MessageDelta(content="test content"),
                    finish_reason=None,
                ),
            ],
        )
        .model_dump_json()
        .encode()
    )

    result = provider._extract_stream_delta(event)

    assert result.delta == "test content"
    assert result.tool_call_requests is None


def test_extract_stream_delta_with_tool_calls() -> None:
    """Test successful extraction of tool calls from stream delta."""
    provider = _TestOpenAIProviderBase()

    event = (
        StreamedResponse(
            choices=[
                ChoiceDelta(
                    index=0,
                    delta=ChoiceDelta.MessageDelta(
                        content=None,
                        tool_calls=[
                            StreamedToolCall(
                                id="call_123",
                                type="function",
                                function=StreamedToolCallFunction(
                                    name="test_tool",
                                    arguments='{"arg1": "value1"}',
                                ),
                                index=0,
                            ),
                        ],
                    ),
                    finish_reason=None,
                ),
            ],
        )
        .model_dump_json()
        .encode()
    )

    result = provider._extract_stream_delta(event)

    assert result.delta is None
    assert result.tool_call_requests is not None
    assert len(result.tool_call_requests) == 1
    tool_call = result.tool_call_requests[0]
    assert tool_call.id == "call_123"
    assert tool_call.tool_name == "test_tool"
    assert tool_call.arguments == '{"arg1": "value1"}'


def test_extract_stream_delta_with_usage() -> None:
    """Test successful extraction of usage information from stream delta."""
    provider = _TestOpenAIProviderBase()

    event = (
        StreamedResponse(
            choices=[
                ChoiceDelta(
                    index=0,
                    delta=ChoiceDelta.MessageDelta(content="test"),
                    finish_reason=None,
                ),
            ],
            usage=Usage(
                prompt_tokens=10,
                completion_tokens=5,
                total_tokens=15,
            ),
        )
        .model_dump_json()
        .encode()
    )

    val = provider._extract_stream_delta(event)

    assert val.usage is not None
    assert val.usage.prompt_token_count == 10
    assert val.usage.completion_token_count == 5


def test_extract_stream_delta_empty_choices() -> None:
    """Test handling of response with empty choices."""
    provider = _TestOpenAIProviderBase()

    event = StreamedResponse(choices=[]).model_dump_json().encode()

    result = provider._extract_stream_delta(event)
    assert result.is_empty()


def test_extract_stream_delta_with_incomplete_tool_call() -> None:
    """Test that incomplete JSON in tool call arguments results in no tool call being added."""
    provider = _TestOpenAIProviderBase()

    # Create an event with a tool call that has incomplete JSON (missing closing brace)
    event = (
        StreamedResponse(
            choices=[
                ChoiceDelta(
                    index=0,
                    delta=ChoiceDelta.MessageDelta(
                        content=None,
                        tool_calls=[
                            StreamedToolCall(
                                id="call_456",
                                type="function",
                                function=StreamedToolCallFunction(
                                    name="incomplete_tool",
                                    arguments='{"arg1": "value1"',
                                ),
                                index=0,
                            ),
                        ],
                    ),
                    finish_reason=None,
                ),
            ],
        )
        .model_dump_json()
        .encode()
    )

    result = provider._extract_stream_delta(event)

    assert result.delta is None
    assert result.tool_call_requests == [
        ToolCallRequestDelta(
            id="call_456",
            idx=0,
            tool_name="incomplete_tool",
            arguments='{"arg1": "value1"',
        ),
    ]


class TestExtractStreamDelta:
    def _load_fixture(self, provider: OpenAIProviderBase[_TestProviderConfig], fixture_name: str):
        raw_completion = RawCompletion(response="", usage=LLMUsage())
        streaming_context = provider._streaming_context(raw_completion)

        fixture_data = fixtures_json(f"openai/{fixture_name}")["events"]
        for sse in fixture_data:
            delta = provider._extract_stream_delta(json.dumps(sse).encode())
            streaming_context.add_chunk(delta)

        final_chunk = streaming_context.complete(
            lambda raw, reasoning, tool_calls: provider._build_structured_output(
                lambda x: x,
                raw,
                reasoning,
                tool_calls,
            ),
        )

        return streaming_context, final_chunk

    def test_extract_stream_delta_from_fixture(self) -> None:
        """Test processing streaming tool call events from a fixture file and verify aggregated tool calls."""

        provider = _TestOpenAIProviderBase()
        _, final_chunk = self._load_fixture(
            provider,
            "stream_with_tools.json",
        )

        assert final_chunk.final_chunk
        assert final_chunk.final_chunk.tool_call_requests == [
            ToolCallRequest(
                index=0,
                tool_name="get_temperature",
                tool_input_dict={"city_code": "125321"},
                id="call_AU0Fw2imWtuWlmaLHXnwkZCQ",
            ),
            ToolCallRequest(
                index=1,
                tool_name="get_rain_probability",
                tool_input_dict={"city_code": "125321"},
                id="call_VtkV4ZpNh3nHxS84JnCiriVj",
            ),
            ToolCallRequest(
                index=2,
                tool_name="get_wind_speed",
                tool_input_dict={"city_code": "125321"},
                id="call_wfkYlsz09dYJvzoXn7spvgpy",
            ),
            ToolCallRequest(
                index=3,
                tool_name="get_weather_conditions",
                tool_input_dict={"city_code": "125321"},
                id="call_3oHEEsy8LFQectXu7wTSOiPx",
            ),
        ]


@pytest.mark.parametrize(
    ("messages", "expected_token_count"),
    [
        (
            # Single simple text message
            [{"role": "user", "content": "Hello, world!"}],
            # 3 (boilerplate) + (4 (message boilerplate) + 4 (tokens in "Hello, world!")) = 11
            11,
        ),
        (
            # Multiple text messages
            [
                {"role": "user", "content": "Hello, world!"},
                {"role": "assistant", "content": "Hi there!"},
            ],
            # 3 (boilerplate) + 2 * (4 (message boilerplate)) + 4 (tokens in "Hello, world!") + 3 (tokens in "Hi there!") = 18
            18,
        ),
        (
            # Empty message list
            [],
            # 3 (boilerplate only)
            3,
        ),
        (
            # Message with empty content
            [{"role": "user", "content": ""}],
            # 3 (boilerplate) + 4 (message boilerplate) + 0 (tokens in empty content) = 7
            7,
        ),
        (
            # Message with tool call
            [{"role": "tool", "tool_call_id": "tc_valid", "content": "Hello"}],
            8,  # 3 (boilerplate) + 4 (message boilerplate) + 1 (tool call)
        ),
        (
            # Message with tool call and other messages
            [
                {"role": "tool", "tool_call_id": "tc_valid", "content": "Hello"},
                {"role": "user", "content": "Hello, world!"},
            ],
            16,  # 3 (boilerplate) + 4 (message boilerplate) + 1 (tool call) + 4 (message boilerplate) + 4 (tokens in "Hello, world!")
        ),
    ],
)
def test_compute_prompt_token_count(messages: list[dict[str, Any]], expected_token_count: int) -> None:
    """Test token count calculation for different message configurations."""
    provider = _TestOpenAIProviderBase()
    model = Model.GPT_4O_2024_08_06

    result = provider._compute_prompt_token_count(messages, model)
    # This is a high-level smoke test that '_compute_prompt_token_count' does not raise and return a value
    assert result == expected_token_count


class TestBuildRequest:
    async def test_with_tool_calls(self, base_provider: _TestOpenAIProviderBase):
        messages = [
            Message(
                role="system",
                content=[
                    MessageContent(
                        text="Be concise, reply with one sentence.Use the `get_lat_lng` tool to get the latitude and longitude of the locations, then use the `get_weather` tool to get the weather.",
                    ),
                ],
            ),
            Message(
                role="user",
                content=[
                    MessageContent(
                        text="What is the weather like in London?",
                    ),
                ],
            ),
            Message(
                role="assistant",
                content=[
                    MessageContent(
                        tool_call_request=ToolCallRequest(
                            tool_name="get_lat_lng",
                            tool_input_dict={"location_description": "London"},
                            id="call_ucYQgwUMFhWu2e91vA9FgRCj",
                        ),
                    ),
                ],
            ),
            Message(
                role="user",
                content=[
                    MessageContent(
                        tool_call_result=ToolCallResult(
                            tool_name="",
                            tool_input_dict={},
                            id="call_ucYQgwUMFhWu2e91vA9FgRCj",
                            result='{"lat":51.5074456,"lng":-0.1277653}',
                        ),
                    ),
                ],
            ),
        ]

        req = base_provider._build_request(
            [m.to_deprecated() for m in messages],
            ProviderOptions(model=Model.GPT_4O_2024_08_06),
            stream=False,
        )
        assert req.model_dump(exclude_none=True)["messages"] == [
            {
                "role": "system",
                "content": "Be concise, reply with one sentence.Use the `get_lat_lng` tool to get the latitude and longitude of the locations, then use the `get_weather` tool to get the weather.",
            },
            {"role": "user", "content": "What is the weather like in London?"},
            {
                "content": [],
                "role": "assistant",
                "tool_calls": [
                    {
                        "function": {
                            "arguments": '{"location_description": "London"}',
                            "name": "get_lat_lng",
                        },
                        "id": "call_ucYQgwUMFhWu2e91vA9FgRCj",
                        "type": "function",
                    },
                ],
            },
            {
                "content": '{"lat":51.5074456,"lng":-0.1277653}',
                "role": "tool",
                "tool_call_id": "call_ucYQgwUMFhWu2e91vA9FgRCj",
            },
        ]

    async def test_with_unsupported_temperature(self, base_provider: _TestOpenAIProviderBase):
        messages = [
            Message(
                role="system",
                content=[MessageContent(text="Be concise, reply with one sentence.")],
            ),
        ]
        req = cast(
            CompletionRequest,
            base_provider._build_request(
                [m.to_deprecated() for m in messages],
                # O3 mini does not support temperature
                ProviderOptions(model=Model.O3_MINI_2025_01_31, temperature=0.5),
                stream=False,
            ),
        )
        assert req.temperature is None


class TestInvalidRequestError:
    def test_too_many_images(self, base_provider: _TestOpenAIProviderBase):
        """Test that 'Too many images in request' error returns ProviderInvalidFileError."""

        # Create a mock OpenAI error payload
        error_payload = OpenAIError(
            error=OpenAIError.Payload(
                code="invalid_request_error",
                message="Too many images in request. Only 5 images are allowed per request.",
                type="invalid_request_error",
                param="messages",
            ),
        )

        # Create a mock HTTP response
        response = Response(
            status_code=400,
            headers={},
            content=b'{"error": {"code": "invalid_request_error", "message": "Too many images in request. Only 5 images are allowed per request.", "type": "invalid_request_error", "param": "messages"}}',
        )

        # Test the _invalid_request_error method
        result = base_provider._invalid_request_error(error_payload, response)

        # Verify it returns ProviderInvalidFileError
        assert isinstance(result, ProviderInvalidFileError)
        assert "Too many images in request" in str(result)

    def test_does_not_support_content(self, base_provider: _TestOpenAIProviderBase):
        error_payload = OpenAIError(
            error=OpenAIError.Payload(
                code="invalid_request_error",
                message="Invalid Value: 'file'. This model does not support file content types.",
                type="invalid_request_error",
                param="model",
            ),
        )

        # Create a mock HTTP response
        response = Response(status_code=400)

        # Test the _invalid_request_error method
        result = base_provider._invalid_request_error(error_payload, response)

        assert isinstance(result, ModelDoesNotSupportModeError)
        assert "Invalid Value: 'file'. This model does not support file content types." in str(result)


class TestResponseFormat:
    @pytest.fixture
    def base_options(self) -> ProviderOptions:
        return ProviderOptions(
            model=Model.GPT_4O_2024_08_06,
            output_schema={"type": "object"},
            structured_generation=True,
        )

    @pytest.mark.parametrize(
        ("supports_structured_output", "supports_json_mode", "expected_cls"),
        [
            pytest.param(True, True, JSONSchemaResponseFormat, id="structured_output_json_mode"),
            pytest.param(True, False, JSONSchemaResponseFormat, id="json_mode_unsupported"),
            pytest.param(False, True, JSONResponseFormat, id="json_mode"),
            pytest.param(False, False, TextResponseFormat, id="text_mode"),
        ],
    )
    async def test_with_structured_output(
        self,
        base_provider: _TestOpenAIProviderBase,
        base_options: ProviderOptions,
        supports_structured_output: bool,
        supports_json_mode: bool,
        expected_cls: type[BaseModel],
    ):
        format = base_provider._response_format(
            options=base_options,
            supports=test_models.fake_model_data(
                supports_structured_output=supports_structured_output,
                supports_json_mode=supports_json_mode,
            ),
        )
        assert isinstance(format, expected_cls)
