from typing import TYPE_CHECKING, Literal, cast
from unittest.mock import ANY, MagicMock

import pytest
from inline_snapshot import snapshot
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import assert_type

from magentic.chat_model.message import (
    AnyMessage,
    AssistantMessage,
    FunctionResultMessage,
    ImageBytes,
    ImageUrl,
    Placeholder,
    SystemMessage,
    ToolResultMessage,
    Usage,
    UserMessage,
)
from magentic.function_call import FunctionCall

if TYPE_CHECKING:
    from collections.abc import Sequence


def test_placeholder_format():
    class Country(BaseModel):
        name: str

    placeholder = Placeholder(Country, "country")

    assert_type(placeholder, Placeholder[Country])
    assert placeholder.name == "country"
    assert placeholder.format(country=Country(name="USA")) == Country(name="USA")


def test_placeholder_coercion():
    placeholder_str = Placeholder(str, "my_string")
    assert placeholder_str.format(my_string="test") == "test"
    placeholder_list_str = Placeholder(list[str], "my_list_str")
    assert placeholder_list_str.format(my_list_str=["test"]) == ["test"]
    assert placeholder_list_str.format(my_list_str={"test"}) == ["test"]  # set -> list
    with pytest.raises(ValueError):  # noqa: PT011
        placeholder_str.format(my_string=True)  # bool does not coerce to str


@pytest.mark.parametrize(
    ("message", "message_model_dump"),
    [
        (SystemMessage("Hello"), {"role": "system", "content": "Hello"}),
        (UserMessage("Hello"), {"role": "user", "content": "Hello"}),
        (AssistantMessage("Hello"), {"role": "assistant", "content": "Hello"}),
        (AssistantMessage(42), {"role": "assistant", "content": 42}),
        (
            ToolResultMessage(3, "unique_id"),
            {"role": "tool", "content": 3, "tool_call_id": "unique_id"},
        ),
        (
            FunctionResultMessage(3, FunctionCall(MagicMock(), 1, 2)),
            {"role": "tool", "content": 3, "tool_call_id": ANY},
        ),
    ],
)
def test_message_model_dump(message, message_model_dump):
    assert message.model_dump() == message_model_dump


@pytest.mark.parametrize(
    ("message", "message_repr"),
    [
        (SystemMessage("Hello"), "SystemMessage('Hello')"),
        (UserMessage("Hello"), "UserMessage('Hello')"),
        (AssistantMessage("Hello"), "AssistantMessage('Hello')"),
        (AssistantMessage(42), "AssistantMessage(42)"),
        (
            FunctionResultMessage(
                3, FunctionCall(MagicMock(__repr__=lambda x: "plus_repr"), 1, 2)
            ),
            "FunctionResultMessage(3, FunctionCall(plus_repr, 1, 2))",
        ),
    ],
)
def test_message_repr(message, message_repr):
    assert repr(message) == message_repr


def test_image_bytes_jpg(image_bytes_jpg: bytes) -> None:
    image_bytes = ImageBytes(image_bytes_jpg)
    assert image_bytes.mime_type == "image/jpeg"
    assert image_bytes.as_base64() == snapshot(
        "/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_image_bytes_png(image_bytes_png: bytes) -> None:
    image_bytes = ImageBytes(image_bytes_png)
    assert image_bytes.mime_type == "image/png"
    assert image_bytes.as_base64() == snapshot(
        " "
    )


def test_image_bytes_invalid():
    with pytest.raises(ValidationError):
        ImageBytes(b"invalid")


def test_user_message_format():
    user_message = UserMessage("Hello {x}")
    user_message_formatted = user_message.format(x="world")

    assert_type(user_message_formatted, UserMessage[str])
    assert_type(user_message_formatted.content, str)
    assert user_message_formatted == UserMessage("Hello world")


def test_user_message_format_type_hints():
    if TYPE_CHECKING:  # Avoid runtime error for None missing `format` method
        assert_type(cast(UserMessage[Literal["x"]], None).format(), UserMessage[str])
        # mypy does not convert `Literal` to `str` in these cases but pyright does
        # issue: https://github.com/python/mypy/issues/18419
        # assert_type(cast(UserMessage[Sequence[Literal["x"]]], None).format(), UserMessage[Sequence[str]])  # noqa: ERA001
        # assert_type(cast(UserMessage[Sequence[Literal["x"] | ImageBytes]], None).format(), UserMessage[Sequence[str | ImageBytes]])  # noqa: ERA001

        assert_type(cast(UserMessage[str], None).format(), UserMessage[str])

        assert_type(
            cast(UserMessage[Sequence[str]], None).format(), UserMessage[Sequence[str]]
        )
        assert_type(
            cast(UserMessage[str | Sequence[str]], None).format(),
            UserMessage[str | Sequence[str]],
        )

        assert_type(
            cast(UserMessage[Sequence[ImageBytes]], None).format(),
            UserMessage[Sequence[ImageBytes]],
        )
        assert_type(
            cast(UserMessage[Sequence[str | ImageBytes]], None).format(),
            UserMessage[Sequence[str | ImageBytes]],
        )
        assert_type(
            cast(UserMessage[str | Sequence[str | ImageBytes]], None).format(),
            UserMessage[str | Sequence[str | ImageBytes]],
        )

        assert_type(
            cast(UserMessage[Sequence[Placeholder[ImageUrl]]], None).format(),
            UserMessage[Sequence[ImageUrl]],
        )
        assert_type(
            cast(UserMessage[Sequence[str | Placeholder[ImageUrl]]], None).format(),
            UserMessage[Sequence[str | ImageUrl]],
        )
        assert_type(
            cast(
                UserMessage[Sequence[ImageBytes | Placeholder[ImageUrl]]], None
            ).format(),
            UserMessage[Sequence[ImageBytes | ImageUrl]],
        )
        assert_type(
            cast(
                UserMessage[Sequence[str | ImageBytes | Placeholder[ImageUrl]]], None
            ).format(),
            UserMessage[Sequence[str | ImageBytes | ImageUrl]],
        )

        assert_type(
            cast(UserMessage[str | Sequence[Placeholder[ImageUrl]]], None).format(),
            UserMessage[str | Sequence[ImageUrl]],
        )
        assert_type(
            cast(
                UserMessage[str | Sequence[str | Placeholder[ImageUrl]]], None
            ).format(),
            UserMessage[str | Sequence[str | ImageUrl]],
        )
        assert_type(
            cast(
                UserMessage[str | Sequence[ImageBytes | Placeholder[ImageUrl]]], None
            ).format(),
            UserMessage[str | Sequence[ImageBytes | ImageUrl]],
        )
        assert_type(
            cast(
                UserMessage[str | Sequence[str | ImageBytes | Placeholder[ImageUrl]]],
                None,
            ).format(),
            UserMessage[str | Sequence[str | ImageBytes | ImageUrl]],
        )


def test_assistant_message_usage():
    assistant_message = AssistantMessage("Hello")
    assert assistant_message.usage is None
    assistant_message = AssistantMessage._with_usage(
        "Hello", [Usage(input_tokens=1, output_tokens=2)]
    )
    assert assistant_message.usage == Usage(input_tokens=1, output_tokens=2)


def test_assistant_message_format_str():
    assistant_message = AssistantMessage("Hello {x}")
    assistant_message_formatted = assistant_message.format(x="world")

    assert_type(assistant_message_formatted, AssistantMessage[str])
    assert_type(assistant_message_formatted.content, str)
    assert assistant_message_formatted == AssistantMessage("Hello world")


def test_assistant_message_format_placeholder():
    class Country(BaseModel):
        name: str

    assistant_message = AssistantMessage(Placeholder(Country, "country"))
    assistant_message_formatted = assistant_message.format(country=Country(name="USA"))

    assert_type(assistant_message_formatted, AssistantMessage[Country])
    assert_type(assistant_message_formatted.content, Country)
    assert assistant_message_formatted == AssistantMessage(Country(name="USA"))


def test_assistant_message_format_type_hints():
    if TYPE_CHECKING:  # Avoid runtime error for None missing `format` method
        assert_type(cast(AssistantMessage[str], None).format(), AssistantMessage[str])
        assert_type(
            cast(AssistantMessage[Literal["x"]], None).format(), AssistantMessage[str]
        )
        assert_type(
            cast(AssistantMessage[Placeholder[int]], None).format(),
            AssistantMessage[int],
        )
        assert_type(
            cast(AssistantMessage[str | Placeholder[int]], None).format(),
            AssistantMessage[str | int],
        )
        assert_type(
            cast(AssistantMessage[Literal["x"] | Placeholder[int]], None).format(),
            AssistantMessage[str | int],
        )
        # Literal matches NotPlaceholder rather than str, so we cannot show that this will become str instead of Literal
        # assert_type(cast(AssistantMessage[Literal["x"] | int | Placeholder[int]], None).format(), AssistantMessage[str | int])  # noqa: ERA001


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

    func_call = FunctionCall(plus, 1, 2)
    function_result_message = FunctionResultMessage(3, func_call)
    assert function_result_message == function_result_message
    assert function_result_message == FunctionResultMessage(3, func_call)
    # Different unique ids internally => not equal, despite equal FunctionCalls
    assert function_result_message != FunctionResultMessage(3, FunctionCall(plus, 1, 2))
    assert function_result_message != FunctionResultMessage(7, FunctionCall(plus, 3, 4))


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

    func_call = FunctionCall(plus, 1, 2)
    function_result_message = FunctionResultMessage(3, func_call)
    function_result_message_formatted = function_result_message.format(foo="bar")

    assert_type(function_result_message_formatted, FunctionResultMessage[int])
    assert_type(function_result_message_formatted.content, int)
    assert function_result_message_formatted == FunctionResultMessage(3, func_call)


def test_any_message():
    messages = [
        {"role": "system", "content": "Hello"},
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hello"},
        {"role": "tool", "content": 3, "tool_call_id": "unique_id"},
    ]
    assert TypeAdapter(list[AnyMessage]).validate_python(messages) == [
        SystemMessage("Hello"),
        UserMessage("Hello"),
        AssistantMessage("Hello"),
        ToolResultMessage(3, "unique_id"),
    ]
