# pyright: reportPrivateUsage=false

import copy
import json
from collections.abc import Callable
from typing import Any, cast

import pytest
from httpx import Response
from pytest_httpx import HTTPXMock, IteratorStream

from core.domain.file import File
from core.domain.message import Message, MessageContent, MessageDeprecated
from core.domain.models import Model
from core.domain.reasoning_effort import ReasoningEffort
from core.providers._base.abstract_provider import RawCompletion
from core.providers._base.llm_usage import LLMUsage
from core.providers._base.provider_error import (
    ContentModerationError,
    MaxTokensExceededError,
    ModelDoesNotSupportModeError,
    ProviderError,
    ProviderInternalError,
    ProviderInvalidFileError,
    StructuredGenerationError,
    UnknownProviderError,
)
from core.providers._base.provider_options import ProviderOptions
from core.providers.xai.xai_domain import CompletionRequest
from core.providers.xai.xai_provider import XAIConfig, XAIProvider
from tests import fake_models as test_models
from tests.utils import fixture_bytes, fixtures_json, fixtures_stream


@pytest.fixture
def xai_provider():
    return XAIProvider(
        config=XAIConfig(api_key="test"),
    )


def _output_factory(x: str) -> Any:
    return json.loads(x)


class TestBuildRequest:
    def test_build_request(self, xai_provider: XAIProvider):
        request = xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
            messages=[
                MessageDeprecated(role=MessageDeprecated.Role.SYSTEM, content="Hello 1"),
                MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello"),
            ],
            options=ProviderOptions(model=Model.GROK_4_0709, max_tokens=10, temperature=0),
            stream=False,
        )
        assert isinstance(request, CompletionRequest)
        # We can exclude None values because the HTTPxProvider does the same
        assert request.model_dump(include={"messages"}, exclude_none=True)["messages"] == [
            {
                "role": "system",
                "content": "Hello 1",
            },
            {
                "role": "user",
                "content": "Hello",
            },
        ]
        assert request.temperature == 0
        assert request.max_tokens == 10

    def test_build_request_without_max_tokens(self, xai_provider: XAIProvider):
        request = xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
            messages=[
                MessageDeprecated(role=MessageDeprecated.Role.SYSTEM, content="Hello 1"),
                MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello"),
            ],
            options=ProviderOptions(model=Model.GROK_4_0709, temperature=0),
            stream=False,
        )
        assert isinstance(request, CompletionRequest)
        # We can exclude None values because the HTTPxProvider does the same
        assert request.model_dump(include={"messages"}, exclude_none=True)["messages"] == [
            {
                "role": "system",
                "content": "Hello 1",
            },
            {
                "role": "user",
                "content": "Hello",
            },
        ]
        assert request.temperature == 0
        assert request.max_tokens is None
        # model_data = get_model_data(Model.GROK_4_0709)
        # if model_data.max_tokens_data.max_output_tokens:
        #     assert request.max_tokens == model_data.max_tokens_data.max_output_tokens
        # else:
        #     assert request.max_tokens == model_data.max_tokens_data.max_tokens

    def test_build_request_with_reasoning_effort(self, xai_provider: XAIProvider):
        request = cast(
            CompletionRequest,
            xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
                messages=[MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello")],
                options=ProviderOptions(
                    model=Model.GROK_3_MINI_BETA,
                    max_tokens=10,
                    temperature=0,
                    reasoning_effort=ReasoningEffort.LOW,
                ),
                stream=False,
            ),
        )
        # We can exclude None values because the HTTPxProvider does the same
        assert request.model_dump(include={"messages", "reasoning_effort", "model"}, exclude_none=True) == {
            "model": "grok-3-mini-beta",
            "messages": [
                {
                    "role": "user",
                    "content": "Hello",
                },
            ],
            "reasoning_effort": "low",
        }

    def test_build_request_with_reasoning_effort_medium(self, xai_provider: XAIProvider):
        request = cast(
            CompletionRequest,
            xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
                messages=[MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello")],
                options=ProviderOptions(
                    model=Model.GROK_3_MINI_BETA,
                    max_tokens=10,
                    temperature=0,
                    reasoning_effort=ReasoningEffort.MEDIUM,
                ),
                stream=False,
            ),
        )
        # We can exclude None values because the HTTPxProvider does the same
        assert request.model_dump(include={"messages", "reasoning_effort", "model"}, exclude_none=True) == {
            "model": "grok-3-mini-beta",
            "messages": [
                {
                    "role": "user",
                    "content": "Hello",
                },
            ],
        }

    def test_build_request_with_tool_choice(self, xai_provider: XAIProvider):
        request = cast(
            CompletionRequest,
            xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
                messages=[MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello")],
                options=ProviderOptions(
                    model=Model.GROK_4_0709,
                    tool_choice="auto",
                ),
                stream=False,
            ),
        )
        # We can exclude None values because the HTTPxProvider does the same
        assert request.model_dump(include={"messages", "tool_choice", "model"}, exclude_none=True) == {
            "model": "grok-4-0709",
            "messages": [
                {
                    "role": "user",
                    "content": "Hello",
                },
            ],
            "tool_choice": "auto",
        }

    @pytest.mark.parametrize(
        ("reasoning", "expected_model", "expected_reasoning_effort"),
        [
            (ReasoningEffort.DISABLED, "grok-4-fast-non-reasoning", None),
            (0, "grok-4-fast-non-reasoning", None),
            (ReasoningEffort.LOW, "grok-4-fast-reasoning", "low"),
            (ReasoningEffort.MEDIUM, "grok-4-fast-reasoning", None),
        ],
    )
    def test_build_request_grok_4_fast_non_reasoning(
        self,
        xai_provider: XAIProvider,
        reasoning: ReasoningEffort | int | None,
        expected_model: str,
        expected_reasoning_effort: str | None,
    ):
        options = ProviderOptions(model=Model.GROK_4_FAST)
        if isinstance(reasoning, int):
            options.reasoning_budget = reasoning
        elif isinstance(reasoning, ReasoningEffort):
            options.reasoning_effort = reasoning
        request = cast(
            CompletionRequest,
            xai_provider._build_request(  # pyright: ignore [reportPrivateUsage]
                messages=[MessageDeprecated(role=MessageDeprecated.Role.USER, content="Hello")],
                options=options,
                stream=False,
            ),
        )
        assert request.model == expected_model
        assert request.reasoning_effort == expected_reasoning_effort


def mock_xai_stream(httpx_mock: HTTPXMock):
    httpx_mock.add_response(
        url="https://api.x.ai/v1/chat/completions",
        stream=IteratorStream(
            [
                b'data: {"id":"1","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,',
                b'"finish_reason":null}]}\n\ndata: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"{\\n"},"logprobs":null,"finish_reason":null}]}\n\n',
                b'data: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"\\"greeting\\": \\"Hello James!\\"\\n}"},"logprobs":null,"finish_reason":null}]}\n\n',
                b"data: [DONE]\n\n",
            ],
        ),
    )


class TestSingleStream:
    async def test_stream_data(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            stream=IteratorStream(
                [
                    b'data: {"id":"1","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,',
                    b'"finish_reason":null}]}\n\ndata: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"{\\n"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b'data: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"\\"greeting\\": \\"Hello James!\\"\\n}"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b"data: [DONE]\n\n",
                ],
            ),
        )

        provider = XAIProvider()
        raw = RawCompletion(usage=LLMUsage(), response="")

        raw_chunks = provider._single_stream(  # pyright: ignore [reportPrivateUsage]
            request={"messages": [{"role": "user", "content": "Hello"}]},
            output_factory=_output_factory,
            raw_completion=raw,
            options=ProviderOptions(model=Model.GPT_3_5_TURBO_1106, max_tokens=10, temperature=0, output_schema={}),
        )

        parsed_chunks = [o async for o in raw_chunks]

        assert len(parsed_chunks) == 4
        final_chunk = parsed_chunks[-1].final_chunk
        assert final_chunk is not None
        assert final_chunk.agent_output == {"greeting": "Hello James!"}

        assert len(httpx_mock.get_requests()) == 1

    # TODO: check for usage and add tests for error handling

    async def test_stream_data_audios(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            stream=IteratorStream(
                [
                    b'data: {"id":"1","object":"chat.completion.chunk","created":1720404416,"model":"gpt-4o-audio-preview","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,',
                    b'"finish_reason":null}]}\n\ndata: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"gpt-4o-audio-preview","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"{\\n"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b'data: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"gpt-4o-audio-preview","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"\\"answer\\": \\"Oh it has 30 words!\\"\\n}"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b"data: [DONE]\n\n",
                ],
            ),
        )

        provider = XAIProvider()
        raw = RawCompletion(usage=LLMUsage(), response="")

        raw_chunks = provider._single_stream(  # pyright: ignore [reportPrivateUsage]
            request={
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "Can you reply to this audio?"},
                            {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,data234avsrtgsd"}},
                        ],
                    },
                ],
            },
            output_factory=_output_factory,
            raw_completion=raw,
            options=ProviderOptions(
                model=Model.GROK_3_MINI_BETA,
                max_tokens=10,
                temperature=0,
                output_schema={},
            ),
        )

        parsed_chunks = [o async for o in raw_chunks]

        assert len(parsed_chunks) == 4
        final_chunk = parsed_chunks[-1].final_chunk
        assert final_chunk is not None
        assert final_chunk.agent_output == {"answer": "Oh it has 30 words!"}

        assert len(httpx_mock.get_requests()) == 1

    async def test_stream_reasoning(self, httpx_mock: HTTPXMock, xai_provider: XAIProvider):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            stream=IteratorStream(fixtures_stream("xai", "stream_reasoning.txt")),
        )
        streamer = xai_provider.stream(
            [Message.with_text("Hello")],
            options=ProviderOptions(model=Model.GROK_3_MINI_BETA),
            output_factory=_output_factory,
        )

        chunks = [copy.deepcopy(o) async for o in streamer]
        # assert len(chunks) == 9

        reasoning = [c.reasoning for c in chunks]
        assert reasoning[0] == "First"
        assert reasoning[1] == " response"


class TestStream:
    # Tests overlap with single stream above but check the entire structure
    async def test_stream(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            stream=IteratorStream(
                [
                    b'data: {"id":"1","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,',
                    b'"finish_reason":null}]}\n\ndata: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"{\\n"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b'data: {"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","choices":[{"index":0,"delta":{"content":"\\"greeting\\": \\"Hello James!\\"\\n}"},"logprobs":null,"finish_reason":null}]}\n\n',
                    b"data: [DONE]\n\n",
                ],
            ),
        )

        provider = XAIProvider()

        streamer = provider.stream(
            [Message.with_text("Hello")],
            options=ProviderOptions(
                model=Model.GROK_4_FAST,
                max_tokens=10,
                temperature=0,
                output_schema={"type": "object"},
            ),
            output_factory=_output_factory,
        )
        chunks = [o async for o in streamer]
        assert len(chunks) == 4

        # Not sure why the pyright in the CI reports an error here
        request = httpx_mock.get_requests()[0]
        assert request.method == "POST"  # pyright: ignore reportUnknownMemberType
        body = json.loads(request.read().decode())
        assert body == {
            "max_tokens": 10,
            "model": "grok-4-fast-reasoning",
            "messages": [
                {
                    "content": "Hello",
                    "role": "user",
                },
            ],
            # "response_format": {
            #     "type": "json_object",
            # },
            "stream": True,
            "stream_options": {
                "include_usage": True,
            },
            "temperature": 0.0,
            # "store": True,
        }

    async def test_stream_error(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            status_code=400,
            json={"error": "blabla"},
        )

        provider = XAIProvider()

        streamer = provider.stream(
            [Message.with_text("Hello")],
            options=ProviderOptions(model=Model.GPT_3_5_TURBO_1106, max_tokens=10, temperature=0),
            output_factory=_output_factory,
        )
        # TODO: be stricter about what error is returned here
        with pytest.raises(UnknownProviderError) as e:
            [chunk async for chunk in streamer]

        assert e.value.capture
        assert str(e.value) == "blabla"


class TestComplete:
    # Tests overlap with single stream above but check the entire structure
    async def test_complete_images(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "completion.json"),
        )

        provider = XAIProvider()

        o = await provider.complete(
            [
                Message(
                    role="user",
                    content=[
                        MessageContent(text="Hello"),
                        MessageContent(file=File(data="data", content_type="image/png")),
                    ],
                ),
            ],
            options=ProviderOptions(
                model=Model.GROK_4_FAST,
                max_tokens=10,
                temperature=0,
                output_schema={"type": "object"},
            ),
            output_factory=_output_factory,
        )
        assert o.agent_output
        assert o.tool_call_requests is None
        # Not sure why the pyright in the CI reports an error here
        request = httpx_mock.get_requests()[0]
        assert request.method == "POST"  # pyright: ignore reportUnknownMemberType
        body = json.loads(request.read().decode())
        assert body == {
            "max_tokens": 10,
            "model": "grok-4-fast-reasoning",
            "messages": [
                {
                    "content": [
                        {
                            "text": "Hello",
                            "type": "text",
                        },
                        {
                            "image_url": {
                                "url": "data:image/png;base64,data",
                            },
                            "type": "image_url",
                        },
                    ],
                    "role": "user",
                },
            ],
            # "response_format": {
            #     "type": "json_object",
            # },
            "stream": False,
            "temperature": 0.0,
            # "store": True,
        }

    async def test_complete_text_only(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "completion.json"),
        )

        provider = XAIProvider()

        o = await provider.complete(
            [
                Message(
                    role="user",
                    content=[
                        MessageContent(text="Hello"),
                    ],
                ),
            ],
            options=ProviderOptions(
                model=Model.GPT_3_5_TURBO_1106,
                max_tokens=10,
                temperature=0,
                # No output schema means the model will return a string
                output_schema=None,
            ),
            output_factory=lambda x: x,
        )
        assert o.agent_output
        assert o.tool_call_requests is None
        # Not sure why the pyright in the CI reports an error here
        request = httpx_mock.get_requests()[0]
        assert request.method == "POST"  # pyright: ignore reportUnknownMemberType
        body = json.loads(request.read().decode())
        assert "response_format" not in body

    async def test_complete_500(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            status_code=500,
            text="Internal Server Error",
        )

        provider = XAIProvider()

        with pytest.raises(ProviderInternalError) as e:
            await provider.complete(
                [Message.with_text("Hello")],
                options=ProviderOptions(model=Model.GPT_3_5_TURBO_1106, max_tokens=10, temperature=0),
                output_factory=_output_factory,
            )

        details = e.value.serialized().details
        assert details
        assert details.get("provider_error") == {"raw": "Internal Server Error"}

    async def test_complete_structured(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "completion.json"),
        )
        provider = XAIProvider()

        o = await provider.complete(
            [
                Message(
                    role="user",
                    content=[
                        MessageContent(text="Hello"),
                        MessageContent(file=File(data="data", content_type="image/png")),
                    ],
                ),
            ],
            options=ProviderOptions(
                model=Model.GROK_CODE_FAST_1,
                max_tokens=10,
                temperature=0,
                task_name="hello",
                structured_generation=True,
                output_schema={"type": "object"},
            ),
            output_factory=_output_factory,
        )
        assert o.agent_output
        assert o.tool_call_requests is None

        # Not sure why the pyright in the CI reports an error here
        request = httpx_mock.get_requests()[0]
        assert request.method == "POST"  # pyright: ignore reportUnknownMemberType
        body = json.loads(request.read().decode())
        assert body == {
            "max_tokens": 10,
            "model": "grok-code-fast-1",
            "messages": [
                {
                    "content": [
                        {
                            "text": "Hello",
                            "type": "text",
                        },
                        {
                            "image_url": {
                                "url": "data:image/png;base64,data",
                            },
                            "type": "image_url",
                        },
                    ],
                    "role": "user",
                },
            ],
            "response_format": {
                "type": "json_schema",
                "json_schema": {
                    "name": "hello_a7515052692371909964e67c28b18a0a",
                    "strict": True,
                    "schema": {"type": "object"},
                },
            },
            "stream": False,
            "temperature": 0.0,
        }

    async def test_complete_refusal(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "refusal.json"),
        )
        provider = XAIProvider()
        with pytest.raises(ContentModerationError) as e:
            await provider.complete(
                [Message.with_text("Hello")],
                options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
                output_factory=_output_factory,
            )

        response = e.value.serialized()
        assert response.code == "content_moderation"
        assert response.status_code == 400
        assert response.message == "Model refused to generate a response: I'm sorry, I can't assist with that."
        assert e.value.capture is False
        assert e.value.store_task_run is True

    @pytest.mark.skip(reason="Not sure about content moderation for now")
    async def test_complete_content_moderation(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "content_moderation.json"),
            status_code=400,
        )

        provider = XAIProvider()
        with pytest.raises(ContentModerationError) as e:
            await provider.complete(
                [Message.with_text("Hello")],
                options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
                output_factory=_output_factory,
            )

        response = e.value.serialized()
        assert response.code == "content_moderation"
        assert response.status_code == 400
        assert response.message.startswith("Invalid prompt: your prompt was flagged")

    @pytest.mark.skip(reason="Not sure about max message length for now")
    async def test_max_message_length(self, httpx_mock: HTTPXMock, xai_provider: XAIProvider):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            status_code=400,
            json={"error": {"code": "string_above_max_length", "message": "The string is too long"}},
        )

        with pytest.raises(MaxTokensExceededError) as e:
            await xai_provider.complete(
                [Message.with_text("Hello")],
                options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
                output_factory=_output_factory,
            )

        assert e.value.store_task_run is False
        assert len(httpx_mock.get_requests()) == 1

    async def test_reasoning(self, httpx_mock: HTTPXMock, xai_provider: XAIProvider):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "reasoning.json"),
        )
        completion = await xai_provider.complete(
            [Message.with_text("Hello")],
            options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
            output_factory=_output_factory,
        )

        assert completion.reasoning
        assert completion.reasoning.startswith("First, the user has provided")


class TestCheckValid:
    async def test_valid(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json={
                "id": "chatcmpl-91gL0PXUwQajck2pIp284pR9o7yVo",
                "object": "chat.completion",
                "created": 1710188102,
                "model": "gpt-4",
                "choices": [
                    {
                        "finish_reason": "stop",
                        "index": 0,
                        "message": {"role": "assistant", "content": "{}"},
                    },
                ],
                "usage": {"prompt_tokens": 35, "completion_tokens": 109, "total_tokens": 144},
                "system_fingerprint": "fp_8abb16fa4e",
            },
        )

        provider = XAIProvider()
        assert await provider.check_valid()

        request = httpx_mock.get_requests()[0]
        body = json.loads(request.read().decode())
        assert body["messages"] == [{"content": "Respond with an empty json", "role": "user"}]


class TestExtractStreamDelta:
    def test_extract_stream_delta(self, xai_provider: XAIProvider):
        delta = xai_provider._extract_stream_delta(  # pyright: ignore[reportPrivateUsage]
            b'{"id":"chatcmpl-9iY4Gi66tnBpsuuZ20bUxfiJmXYQC","object":"chat.completion.chunk","created":1720404416,"model":"grok-4-fast-reasoning","system_fingerprint":"fp_44132a4de3","usage": {"prompt_tokens": 35, "completion_tokens": 109, "total_tokens": 144},"choices":[{"index":0,"delta":{"content":"\\"greeting\\": \\"Hello James!\\"\\n}"},"logprobs":null,"finish_reason":null}]}',
        )
        assert delta.delta == '"greeting": "Hello James!"\n}'
        assert delta.usage == LLMUsage(prompt_token_count=35, completion_token_count=109)

    def test_done(self, xai_provider: XAIProvider):
        delta = xai_provider._extract_stream_delta(b"[DONE]")  # pyright: ignore[reportPrivateUsage]
        assert delta.is_empty()


class TestMaxTokensExceededError:
    async def test_max_tokens_exceeded_error(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            json=fixtures_json("xai", "finish_reason_length_completion.json"),
        )

        provider = XAIProvider()
        with pytest.raises(MaxTokensExceededError) as e:
            await provider.complete(
                [Message.with_text("Hello")],
                options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
                output_factory=_output_factory,
            )
        assert (
            e.value.args[0]
            == "Model returned a response with a length finish reason, meaning the maximum number of tokens was exceeded."
        )
        assert e.value.store_task_run is True
        assert len(httpx_mock.get_requests()) == 1

    async def test_max_tokens_exceeded_error_stream(self, httpx_mock: HTTPXMock):
        httpx_mock.add_response(
            url="https://api.x.ai/v1/chat/completions",
            stream=IteratorStream(
                [
                    fixture_bytes("xai", "finish_reason_length_stream_completion.txt"),
                ],
            ),
        )
        provider = XAIProvider()
        with pytest.raises(MaxTokensExceededError) as e:
            async for _ in provider._single_stream(  # pyright: ignore reportPrivateUsage
                {"messages": [{"role": "user", "content": "Hello"}]},
                options=ProviderOptions(model=Model.GPT_4O_2024_08_06, max_tokens=10, temperature=0),
                output_factory=_output_factory,
                raw_completion=RawCompletion(response="", usage=LLMUsage()),
            ):
                pass

        assert (
            e.value.args[0]
            == "Model returned a response with a length finish reason, meaning the maximum number of tokens was exceeded."
        )
        assert e.value.store_task_run is True
        assert len(httpx_mock.get_requests()) == 1


class TestPrepareCompletion:
    async def test_role_before_content(self, xai_provider: XAIProvider):
        """Test that the 'role' key appears before 'content' in the prepared request."""
        request, _ = await xai_provider._prepare_completion(  # pyright: ignore[reportPrivateUsage]
            messages=[Message.with_text("Hello")],
            options=ProviderOptions(model=Model.GPT_3_5_TURBO_1106, max_tokens=10, temperature=0),
            stream=False,
        )

        # Get the first message from the request
        message = request["messages"][0]

        # Get the actual order of keys in the message dictionary
        keys = list(message.keys())

        # Find the indices of 'role' and 'content' in the keys list
        role_index = keys.index("role")
        content_index = keys.index("content")

        assert role_index < content_index, (
            "The 'role' key must appear before the 'content' key in the message dictionary"
        )


class TestUnknownError:
    @pytest.fixture
    def unknown_error_fn(self, xai_provider: XAIProvider):
        # Wrapper to avoid having to silence the private warning
        # and instantiate the response
        def _build_unknown_error(payload: str | dict[str, Any], status_code: int = 400):
            if isinstance(payload, dict):
                payload = json.dumps(payload)
            res = Response(status_code=status_code, text=payload)
            return xai_provider._unknown_error(res)  # pyright: ignore[reportPrivateUsage]

        return _build_unknown_error

    def test_max_tokens_error(self, unknown_error_fn: Callable[[dict[str, Any]], ProviderError]):
        payload = {
            "code": "Client specified an invalid argument",
            "error": "This model's maximum prompt length is 131072 but the request contains 144543 tokens.",
        }
        assert isinstance(unknown_error_fn(payload), MaxTokensExceededError)

    def test_model_does_not_support_mode(self, unknown_error_fn: Callable[[dict[str, Any], int], ProviderError]):
        payload = {
            "code": "The system is not in a state required for the operation's execution",
            "error": 'Unsupported content-type encountered when downloading image. The only supported content types are ["image/jpeg", "image/jpg", "image/png"]',
        }
        assert isinstance(unknown_error_fn(payload, 412), ModelDoesNotSupportModeError)

    def test_wrong_content(self, unknown_error_fn: Callable[[dict[str, Any]], ProviderError]):
        payload = {
            "code": "Client specified an invalid argument",
            "error": "Downloaded response does not contain a valid JPG or PNG image",
        }
        assert isinstance(unknown_error_fn(payload), ProviderInvalidFileError)

    def test_structured_generation_error(self, unknown_error_fn: Callable[[dict[str, Any]], ProviderError]):
        payload = {
            "code": "Client specified an invalid argument",
            "error": "Invalid request content: The model does not support formatted output but some have been specified in the request.",
        }
        assert isinstance(unknown_error_fn(payload), StructuredGenerationError)


class TestResponseFormat:
    def test_response_format_structured_output_disabled(self, xai_provider: XAIProvider):
        options = ProviderOptions(
            model=Model.GPT_4O_2024_08_06,
            max_tokens=10,
            temperature=0,
            structured_generation=False,
            output_schema={"type": "object"},
        )
        model_data = test_models.fake_model_data(supports_structured_output=True)
        response_format = xai_provider._response_format(options, model_data)
        assert response_format is None

    def test_response_format_structured_output_enabled(self, xai_provider: XAIProvider):
        options = ProviderOptions(
            model=Model.GPT_4O_2024_08_06,
            max_tokens=10,
            temperature=0,
            structured_generation=True,
            output_schema={"type": "object"},
        )
        model_data = test_models.fake_model_data(supports_structured_output=True)
        response_format = xai_provider._response_format(options, model_data)
        assert response_format is not None
        assert response_format.type == "json_schema"
        assert response_format.json_schema.json_schema == {"type": "object"}
