import os
from typing import Annotated, Any
from unittest.mock import ANY

import openai
import pytest
from inline_snapshot import snapshot
from openai.types.chat import ChatCompletionMessageParam
from pydantic import AfterValidator, BaseModel

from magentic._pydantic import ConfigDict, with_config
from magentic._streamed_response import AsyncStreamedResponse, StreamedResponse
from magentic.chat_model.base import ToolSchemaParseError
from magentic.chat_model.message import (
    AssistantMessage,
    FunctionResultMessage,
    ImageBytes,
    ImageUrl,
    Message,
    SystemMessage,
    Usage,
    UserMessage,
    _RawMessage,
)
from magentic.chat_model.openai_chat_model import (
    OpenaiChatModel,
    async_message_to_openai_message,
    message_to_openai_message,
)
from magentic.function_call import (
    AsyncParallelFunctionCall,
    FunctionCall,
    ParallelFunctionCall,
)
from magentic.streaming import AsyncStreamedStr, StreamedStr, async_iter


def plus(a: int, b: int) -> int:
    return a + b


message_to_openai_message_test_cases = [
    (
        _RawMessage({"role": "user", "content": "Hello"}),
        {"role": "user", "content": "Hello"},
    ),
    (SystemMessage("Hello"), {"role": "system", "content": "Hello"}),
    (UserMessage("Hello"), {"role": "user", "content": "Hello"}),
    (
        UserMessage([ImageUrl("https://example.com/image.jpg")]),
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": "https://example.com/image.jpg"},
                }
            ],
        },
    ),
    (AssistantMessage("Hello"), {"role": "assistant", "content": "Hello"}),
    (
        AssistantMessage(42),
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "return_int", "arguments": '{"value":42}'},
                }
            ],
        },
    ),
    (
        AssistantMessage(FunctionCall(plus, 1, 2)),
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
                }
            ],
        },
    ),
    (
        AssistantMessage(
            ParallelFunctionCall([FunctionCall(plus, 1, 2), FunctionCall(plus, 3, 4)])
        ),
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
                },
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":3,"b":4}'},
                },
            ],
        },
    ),
    (
        AssistantMessage(
            StreamedResponse([StreamedStr(["Hello"]), FunctionCall(plus, 1, 2)])
        ),
        {
            "role": "assistant",
            "content": "Hello",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
                },
            ],
        },
    ),
    (
        FunctionResultMessage(3, FunctionCall(plus, 1, 2)),
        {
            "role": "tool",
            "tool_call_id": ANY,
            "content": '{"value":3}',
        },
    ),
]


@pytest.mark.parametrize(
    ("message", "expected_openai_message"), message_to_openai_message_test_cases
)
def test_message_to_openai_message(message, expected_openai_message):
    assert message_to_openai_message(message) == expected_openai_message


async_message_to_openai_message_test_cases = [
    *message_to_openai_message_test_cases,
    (
        AssistantMessage(
            AsyncParallelFunctionCall(
                async_iter([FunctionCall(plus, 1, 2), FunctionCall(plus, 3, 4)])
            )
        ),
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
                },
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":3,"b":4}'},
                },
            ],
        },
    ),
    (
        AssistantMessage(
            AsyncStreamedResponse(
                async_iter(
                    [
                        AsyncStreamedStr(async_iter(["Hello", " World"])),
                        FunctionCall(plus, 1, 2),
                    ]
                )
            )
        ),
        {
            "role": "assistant",
            "content": "Hello World",
            "tool_calls": [
                {
                    "id": ANY,
                    "type": "function",
                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
                },
            ],
        },
    ),
]


@pytest.mark.parametrize(
    ("message", "expected_openai_message"), async_message_to_openai_message_test_cases
)
async def test_async_message_to_openai_message(message, expected_openai_message):
    assert await async_message_to_openai_message(message) == expected_openai_message


def test_message_to_openai_message_user_image_message_bytes_jpg(image_bytes_jpg):
    image_message = UserMessage([ImageBytes(image_bytes_jpg)])
    assert message_to_openai_message(image_message) == snapshot(
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": "data:image/jpeg;base64,/9j/4QDoRXhpZgAATU0AKgAAAAgABwESAAMAAAABAAEAAAEaAAUAAAABAAAAYgEbAAUAAAABAAAAagEoAAMAAAABAAIAAAExAAIAAAARAAAAcgITAAMAAAABAAEAAIdpAAQAAAABAAAAhAAAAAAAAABIAAAAAQAAAEgAAAABd3d3Lmlua3NjYXBlLm9yZwAAAAeQAAAHAAAABDAyMjGRAQAHAAAABAECAwCgAAAHAAAABDAxMDCgAQADAAAAAQABAACgAgAEAAAAAQAAADKgAwAEAAAAAQAAAEGkBgADAAAAAQAAAAAAAAAAAAD/2wCEAAEBAQEBAQIBAQIDAgICAwQDAwMDBAYEBAQEBAYHBgYGBgYGBwcHBwcHBwcICAgICAgJCQkJCQsLCwsLCwsLCwsBAgICAwMDBQMDBQsIBggLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLCwsLC//dAAQABP/AABEIAEEAMgMBIgACEQEDEQH/xAGiAAABBQEBAQEBAQAAAAAAAAAAAQIDBAUGBwgJCgsQAAIBAwMCBAMFBQQEAAABfQECAwAEEQUSITFBBhNRYQcicRQygZGhCCNCscEVUtHwJDNicoIJChYXGBkaJSYnKCkqNDU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6g4SFhoeIiYqSk5SVlpeYmZqio6Slpqeoqaqys7S1tre4ubrCw8TFxsfIycrS09TV1tfY2drh4uPk5ebn6Onq8fLz9PX29/j5+gEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoLEQACAQIEBAMEBwUEBAABAncAAQIDEQQFITEGEkFRB2FxEyIygQgUQpGhscEJIzNS8BVictEKFiQ04SXxFxgZGiYnKCkqNTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqCg4SFhoeIiYqSk5SVlpeYmZqio6Slpqeoqaqys7S1tre4ubrCw8TFxsfIycrS09TV1tfY2dri4+Tl5ufo6ery8/T19vf4+fr/2gAMAwEAAhEDEQA/AP6RNf8AEvh/wrZLqPiS8hsYGcRI0zhAzt0Rc9WOOFHJ7CsbRfiP4E8RQ2M+h6rbXK6lLLb2wRxl5YFLSJjqGQKdykAjHNeSftB3t7p2r/Du90+zkv5o/FKlbeJkR3/4l99wpkZUzjplgPcV5npXw7+KGm+KYfi7JocjSSeKLvVG0RJ7f7RFaXFgLIMGMggMu5BI6CXbhuDkYIB7z4n+N/grw14j0Pw4Z1uX1jVJ9JZozkW81vA8zBwAeflVcf7YPQVzvwf+N/h/xhZf2P4l1iwXX31XVrOOzWRUlKWd3PHGuzOd3kxq2OpALYxXi/gHQPFniHV4/HNjpj/8Srx/qd7cWnmxeasMllJZkg7/ACyY5HG4B+gOMkAHr9N+FHja3+H2g6I9kEvLLxrcaxKBJHlLSTULiYSA5xkwyL8o+bBxjqKAPoDSvin8Odb1yfw1pOs2k99b+aHiWQf8sOJAD0PlnhwudnfFHhv4pfDvxfrD+H/DOsW17eJEZ/KifJaIHaXTs6A4BZMgEgelfKdt8JfiZc+C/Cnw0m0p7Y+D1vWk1MzwlL4ta3FtGsW1zJmczCSXzUjAwRzxXqPhv4c+K9P8QfDDUbq0Cr4c8P3VhqDb0JimlitVVOD8wLRNyuRx9KAPpmiiigD/0P6SdU0HR9ansrnVbdJ5NOn+1WzMOYpgjR7l99jsv0NWNT02z1jT5dLv1LQTrscKxQkH3Ugj8DV+vzj/AGwv25NT/Z+8VwfD3wHpltqGqeQtxczXm4wxLJnYgRChLEDJO4ADHXt25fl1fG1lQw8by+7RHFj8woYKj7bEO0djrPip8f8Axl8IfDXjH/hR3gS1vvDfwx+wRa1eSXcdrFbSakR5Kx24/eSZZhuZe+cjvXyH4G/4K3TvrtvafEvwrFBp0jBZbjT5mMkQP8XluPmA7gMD6elfBPxE/aGtviT4ov8Axn4w8F6Lc3+psj3LrLqEauUUKvyJeKvAA6CvlrWbiG9v5rq0t0tI5pGZIIixSMMchF3lmwo4G5icdSa/SsNwthI0fZ4ij71lrfyV9pPre2iVrKx+cYnijFOtz4er7va3S+n2V0st36n9ocMsU8SzwkMjgMpHQg9KkrB0HzbXRLG3nXayQRKwPBBCit6vwbh/irLc69usvnd0Zckk1azX6O2jXY/Y8dleIwapuvG3MroKKKK+iOA//9H+mCv51/8Ago9oGp6T+01qOqXqsINUtLSe3Yj5SiRCI4+jRmv6KK8z+I3wb+F/xcgt7f4kaJbasLXPktMvzx7uoVlwwBwMgHHFe9w7m8ctxft5xvFpp2+X+R4XEOUSzHC+whKzTTXbt+TPiH9mP/gnb8A9T+DGheMfilp8muarrdnFfPuuJYY4VnUOsaLC6fdUgEnOTnGBgD6M0D9hf9lPwXr1t4l0HwhAt5ZuJIWmmnuFV16HZLIyZHbjivrXw3pWmaB4YsdA0WJYLSwhS3hiXokcQ2qo+gAFbkGg29/o+oatLfw272QQrBI2JJt5x8g77e9fxZxf4kcX4zP8xwkM1q04c9VKCk4pQi5aJK1rRWnV9Ls/Ycj4SyTDYDDTeDpykow15U3ey1u13+70PPbkZO0fSr9N2jO4jmnV+u+FXh9iOGKGJeLqxnUrOPw3slFO26WvvO+nY8TifPqeZTpKlG0YJ72627dNEFFFFfrB8sf/0v6L/GHjbTfBq6fDdQzXd3q1ybKxtbdVMk84hknKguURcRQyNl2VcLjOcCvM0/aR+GsmlprURu2to7YXl+wgP/EugMskG65H8O2WGVGC7iPLZsbFLDS+L3grxP4w1zwRP4anlshpGuSXtzdQ+UXghOm30AYLMrK2ZJo0ICk4bIxjcOQP7LXgyLTpdHsNT1CC01GzNjrEYMLf2nC1xNct55MR2l5LifcYvLysrDA+Xb1wjQ5VzvX/AIf9LW/yOScq3M+Raf8AAX9P8Do5f2lPhjYXaWNxczwzPHqzlDEQU/sV2S4RsdGPlu0Y/jVGI4FRv+018LEvtXs9QvJIBosN7NO7JuDDTmWO6CKhZ90TsEIKjcc7NwBIz9c/Zh+HGv8AiLUvE1492txqerWOsOEkARJLGMxeWg28Rzo8qzL/ABiV+Rmp7v8AZ38P3GieI/CkOqXsOkeJDeyS2irBiGXUXMs7RyGLzMGRmcKzEAsf4dqryzy3K51Y4iVJOore84q69Ha+n6GixGOjF01L3e3T7rmndfH7wXpPiCw8KeJILrStRvRalre5EW+D7dO1va+Zskf/AF0iEKE37B/rNlc/onxqm8afFfQdD8MQXUWhX1nq8n2ieJFivHsZbeNXhOTJsBd8ZVA4IIyMGu91H4U6fc/EQ/ErTL+4sLy4gtra8jjSGSK4itHd4siWNyjDzHUtGVJU4PIUrn+B/groPgbUbK9tL67u49Jju4NNgnMYjtLe8dHeJNkaMyr5aqm8sVUAe9dSdBK/W36f57eXmZ2rXt0/4b9PL8D2SiiiuQ6z/9P+mCiiigAooooAKKKKACiiigD/2Q=="
                    },
                }
            ],
        }
    )


def test_message_to_openai_message_user_image_message_bytes_png(image_bytes_png):
    image_message = UserMessage([ImageBytes(image_bytes_png)])
    assert message_to_openai_message(image_message) == snapshot(
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": "data:image/png;base64, "
                    },
                }
            ],
        }
    )


def test_message_to_openai_message_raises():
    class CustomMessage(Message[str]):
        def __init__(self, content: str, **data: Any):
            super().__init__(content, **data)

        def format(self, **kwargs):
            del kwargs
            return CustomMessage(self.content)

    with pytest.raises(NotImplementedError):
        message_to_openai_message(CustomMessage("Hello"))

    @message_to_openai_message.register
    def _(message: CustomMessage) -> ChatCompletionMessageParam:
        return {"role": "user", "content": message.content}

    assert message_to_openai_message(CustomMessage("Hello")) == {
        "role": "user",
        "content": "Hello",
    }


@pytest.mark.openai
def test_openai_chat_model_api_key(monkeypatch):
    openai_api_key = os.environ["OPENAI_API_KEY"]
    monkeypatch.delenv("OPENAI_API_KEY")

    with pytest.raises(openai.OpenAIError):
        chat_model = OpenaiChatModel("gpt-4o")

    chat_model = OpenaiChatModel("gpt-4o", api_key=openai_api_key)
    message = chat_model.complete(messages=[UserMessage("Say hello!")])
    assert isinstance(message.content, str)


@pytest.mark.openai
def test_openai_chat_model_complete_base_url():
    chat_model = OpenaiChatModel("gpt-4o", base_url="https://api.openai.com/v1")
    message = chat_model.complete(messages=[UserMessage("Say hello!")])
    assert isinstance(message.content, str)


@pytest.mark.openai
def test_openai_chat_model_complete_seed():
    chat_model = OpenaiChatModel("gpt-4o", seed=42)
    message1 = chat_model.complete(messages=[UserMessage("Say hello!")])
    message2 = chat_model.complete(messages=[UserMessage("Say hello!")])
    assert message1.content == message2.content


@pytest.mark.openai
def test_openai_chat_model_complete_streamed_response():
    def get_weather(location: str) -> None:
        """Get the weather for a location."""

    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[UserMessage("Tell me your favorite city. Then get its weather.")],
        functions=[get_weather],
        output_types=[StreamedResponse],
    )
    assert isinstance(message.content, StreamedResponse)
    response_items = list(message.content)
    assert len(response_items) == 2
    streamed_str, function_call = response_items
    assert isinstance(streamed_str, StreamedStr)
    assert len(streamed_str.to_string()) > 1  # Check StreamedStr was cached
    assert isinstance(function_call, FunctionCall)
    assert function_call() is None  # Check FunctionCall is successfully called


@pytest.mark.openai
def test_openai_chat_model_complete_pydantic_model_openai_strict():
    class CapitalCity(BaseModel):
        model_config = ConfigDict(openai_strict=True)
        capital: str
        country: str

    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[UserMessage("What is the capital of Ireland?")],
        output_types=[CapitalCity],
    )
    assert isinstance(message.content, CapitalCity)


@pytest.mark.openai
def test_openai_chat_model_complete_function_call_openai_strict():
    @with_config(ConfigDict(openai_strict=True))
    def plus(a: int, b: int) -> int:
        """Sum two numbers."""
        return a + b

    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[UserMessage("Use the tool to sum 1 and 2")],
        functions=[plus],
        output_types=[FunctionCall[int]],
    )
    assert isinstance(message.content, FunctionCall)


@pytest.mark.openai
def test_openai_chat_model_complete_usage():
    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[UserMessage("Say hello!")], output_types=[StreamedStr]
    )
    str(message.content)  # Finish the stream
    assert isinstance(message.usage, Usage)
    assert message.usage.input_tokens > 0
    assert message.usage.output_tokens > 0


@pytest.mark.openai
def test_openai_chat_model_complete_usage_structured_output():
    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[UserMessage("Count to 5")], output_types=[list[int]]
    )
    assert isinstance(message.usage, Usage)
    assert message.usage.input_tokens > 0
    assert message.usage.output_tokens > 0


@pytest.mark.openai
def test_openai_chat_model_complete_no_structured_output_error():
    chat_model = OpenaiChatModel("gpt-4o")
    # Should not raise StructuredOutputError because forced to make tool call
    message: Message[int | bool] = chat_model.complete(
        messages=[
            UserMessage("Tell me a short joke. Return a string, not a tool call.")
        ],
        output_types=[int, bool],
    )
    assert isinstance(message.content, int | bool)


@pytest.mark.openai
def test_openai_chat_model_complete_raises_tool_schema_parse_error():
    def raise_error(v):
        raise ValueError(v)

    class Test(BaseModel):
        value: Annotated[int, AfterValidator(raise_error)]

    chat_model = OpenaiChatModel("gpt-4o")
    with pytest.raises(ToolSchemaParseError):
        chat_model.complete(
            messages=[UserMessage("Return a test value of 42.")],
            output_types=[Test],
        )


@pytest.mark.openai
def test_openai_chat_model_complete_image_bytes(image_bytes_jpg):
    chat_model = OpenaiChatModel("gpt-4o")
    message = chat_model.complete(
        messages=[
            UserMessage(
                ("Describe this image in one word.", ImageBytes(image_bytes_jpg))
            )
        ]
    )
    assert isinstance(message.content, str)


@pytest.mark.openai
async def test_openai_chat_model_acomplete_async_streamed_response():
    def get_weather(location: str) -> None:
        """Get the weather for a location."""

    chat_model = OpenaiChatModel("gpt-4o")
    message = await chat_model.acomplete(
        messages=[UserMessage("Tell me your favorite city. Then get its weather.")],
        functions=[get_weather],
        output_types=[AsyncStreamedResponse],
    )
    assert isinstance(message.content, AsyncStreamedResponse)
    response_items = [x async for x in message.content]
    assert len(response_items) == 2
    streamed_str, function_call = response_items
    assert isinstance(streamed_str, AsyncStreamedStr)
    assert len(await streamed_str.to_string()) > 1  # Check AsyncStreamedStr was cached
    assert isinstance(function_call, FunctionCall)
    assert function_call() is None  # Check FunctionCall is successfully called


@pytest.mark.openai
async def test_openai_chat_model_acomplete_usage():
    chat_model = OpenaiChatModel("gpt-4o")
    message = await chat_model.acomplete(
        messages=[UserMessage("Say hello!")], output_types=[AsyncStreamedStr]
    )
    await message.content.to_string()  # Finish the stream
    assert isinstance(message.usage, Usage)
    assert message.usage.input_tokens > 0
    assert message.usage.output_tokens > 0


@pytest.mark.openai
async def test_openai_chat_model_acomplete_usage_structured_output():
    chat_model = OpenaiChatModel("gpt-4o")
    message = await chat_model.acomplete(
        messages=[UserMessage("Count to 5")], output_types=[list[int]]
    )
    assert isinstance(message.usage, Usage)
    assert message.usage.input_tokens > 0
    assert message.usage.output_tokens > 0


@pytest.mark.openai
async def test_openai_chat_model_acomplete_raises_tool_schema_parse_error():
    def raise_error(v):
        raise ValueError(v)

    class Test(BaseModel):
        value: Annotated[int, AfterValidator(raise_error)]

    chat_model = OpenaiChatModel("gpt-4o")
    with pytest.raises(ToolSchemaParseError):
        await chat_model.acomplete(
            messages=[UserMessage("Return a test value of 42.")],
            output_types=[Test],
        )


def test_openai_chat_model_azure_omits_stream_options(monkeypatch):
    monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test")
    monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "test")
    monkeypatch.setenv("OPENAI_API_VERSION", "test")
    chat_model = OpenaiChatModel("gpt-4o", api_type="azure")
    assert chat_model._get_stream_options() == openai.NOT_GIVEN
