# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import typing
from collections import deque
from types import UnionType
from typing import Any, Deque, Dict, FrozenSet, List, Optional, Set, Tuple, Union

import pytest

from haystack.dataclasses import Answer, ByteStream, ChatMessage, Document
from haystack.utils.type_serialization import (
    _build_pep604_union_type,
    _is_union_type,
    _parse_pep604_union_args,
    deserialize_type,
    serialize_type,
)

TYPING_AND_TYPE_TESTS = [
    # dict
    pytest.param("dict", dict),
    pytest.param("dict[str, int]", dict[str, int]),
    pytest.param("dict[int, str]", dict[int, str]),
    pytest.param("dict[dict, dict]", dict[dict, dict]),
    pytest.param("dict[float, float]", dict[float, float]),
    pytest.param("dict[bool, bool]", dict[bool, bool]),
    # typing Dict
    pytest.param("typing.Dict", Dict),
    pytest.param("typing.Dict[str, int]", Dict[str, int]),
    pytest.param("typing.Dict[int, str]", Dict[int, str]),
    pytest.param("typing.Dict[dict, dict]", Dict[dict, dict]),
    pytest.param("typing.Dict[float, float]", Dict[float, float]),
    pytest.param("typing.Dict[bool, bool]", Dict[bool, bool]),
    # list
    pytest.param("list", list),
    pytest.param("list[int]", list[int]),
    pytest.param("list[str]", list[str]),
    pytest.param("list[dict]", list[dict]),
    pytest.param("list[float]", list[float]),
    pytest.param("list[bool]", list[bool]),
    # typing List
    pytest.param("typing.List", List),
    pytest.param("typing.List[int]", List[int]),
    pytest.param("typing.List[str]", List[str]),
    pytest.param("typing.List[dict]", List[dict]),
    pytest.param("typing.List[float]", List[float]),
    pytest.param("typing.List[bool]", List[bool]),
    # typing Optional
    pytest.param("typing.Optional", Optional),
    pytest.param("typing.Optional[str]", Optional[str]),
    pytest.param("typing.Optional[int]", Optional[int]),
    pytest.param("typing.Optional[dict]", Optional[dict]),
    pytest.param("typing.Optional[float]", Optional[float]),
    pytest.param("typing.Optional[bool]", Optional[bool]),
    # PEP 604 X | None
    pytest.param("str | None", str | None),
    pytest.param("int | None", int | None),
    pytest.param("dict | None", dict | None),
    pytest.param("float | None", float | None),
    pytest.param("bool | None", bool | None),
    pytest.param("list[str] | None", list[str] | None),
    pytest.param("dict[str, int] | None", dict[str, int] | None),
    # set
    pytest.param("set", set),
    pytest.param("set[int]", set[int]),
    pytest.param("set[str]", set[str]),
    pytest.param("set[dict]", set[dict]),
    pytest.param("set[float]", set[float]),
    pytest.param("set[bool]", set[bool]),
    # typing Set
    pytest.param("typing.Set", Set),
    pytest.param("typing.Set[int]", Set[int]),
    pytest.param("typing.Set[str]", Set[str]),
    pytest.param("typing.Set[dict]", Set[dict]),
    pytest.param("typing.Set[float]", Set[float]),
    pytest.param("typing.Set[bool]", Set[bool]),
    # tuple
    pytest.param("tuple", tuple),
    pytest.param("tuple[int]", tuple[int]),
    pytest.param("tuple[str]", tuple[str]),
    pytest.param("tuple[dict]", tuple[dict]),
    pytest.param("tuple[float]", tuple[float]),
    pytest.param("tuple[bool]", tuple[bool]),
    # typing Tuple
    pytest.param("typing.Tuple", Tuple),
    pytest.param("typing.Tuple[int]", Tuple[int]),
    pytest.param("typing.Tuple[str]", Tuple[str]),
    pytest.param("typing.Tuple[dict]", Tuple[dict]),
    pytest.param("typing.Tuple[float]", Tuple[float]),
    pytest.param("typing.Tuple[bool]", Tuple[bool]),
    # Union
    pytest.param("typing.Union", Union),
    pytest.param("typing.Union[str, int]", Union[str, int]),
    pytest.param("typing.Union[int, float]", Union[int, float]),
    pytest.param("typing.Union[dict, str]", Union[dict, str]),
    pytest.param("typing.Union[float, bool]", Union[float, bool]),
    pytest.param("typing.Optional[str]", typing.Union[None, str]),  # Union with None becomes Optional
    # PEP 604 X | Y
    pytest.param("str | int", str | int),
    pytest.param("int | float", int | float),
    pytest.param("dict | str", dict | str),
    pytest.param("float | bool", float | bool),
    pytest.param("str | int | float", str | int | float),
    pytest.param("list[str] | list[int]", list[str] | list[int]),
    pytest.param("dict[str, int] | list[str]", dict[str, int] | list[str]),
    # other
    pytest.param("frozenset", frozenset),
    pytest.param("frozenset[int]", frozenset[int]),
    pytest.param("collections.deque", deque),
    pytest.param("collections.deque[str]", deque[str]),
    # typing Other
    pytest.param("typing.Any", Any),
    pytest.param("typing.FrozenSet", FrozenSet),
    pytest.param("typing.FrozenSet[int]", FrozenSet[int]),
    pytest.param("typing.Deque", Deque),
    pytest.param("typing.Deque[str]", Deque[str]),
]


@pytest.mark.parametrize("output_str, input_type", TYPING_AND_TYPE_TESTS)
def test_output_type_serialization_typing_and_type(output_str, input_type):
    assert serialize_type(input_type) == output_str


@pytest.mark.parametrize("input_str, expected_output", TYPING_AND_TYPE_TESTS)
def test_output_type_deserialization_typing_and_type(input_str, expected_output):
    assert deserialize_type(input_str) == expected_output


def test_output_type_deserialization_typing_no_module():
    assert deserialize_type("List[int]") == List[int]
    assert deserialize_type("Dict[str, int]") == Dict[str, int]
    assert deserialize_type("Set[int]") == Set[int]
    assert deserialize_type("Tuple[int]") == Tuple[int]
    assert deserialize_type("FrozenSet[int]") == FrozenSet[int]
    assert deserialize_type("Deque[str]") == Deque[str]
    assert deserialize_type("Optional[int]") == Optional[int]
    assert deserialize_type("Union[str, int]") == Union[str, int]


def test_output_type_serialization():
    assert serialize_type(str) == "str"
    assert serialize_type(int) == "int"
    assert serialize_type(dict) == "dict"
    assert serialize_type(float) == "float"
    assert serialize_type(bool) == "bool"
    assert serialize_type(None) == "None"


def test_output_type_serialization_string():
    assert serialize_type("str") == "str"
    assert serialize_type("builtins.str") == "builtins.str"


def test_output_type_deserialization():
    assert deserialize_type("str") == str
    assert deserialize_type("int") == int
    assert deserialize_type("dict") == dict
    assert deserialize_type("float") == float
    assert deserialize_type("bool") == bool
    assert deserialize_type("None") is None
    assert deserialize_type("NoneType") == type(None)  # type: ignore


def test_output_builtin_type_deserialization():
    assert deserialize_type("builtins.str") == str
    assert deserialize_type("builtins.int") == int
    assert deserialize_type("builtins.dict") == dict
    assert deserialize_type("builtins.float") == float
    assert deserialize_type("builtins.bool") == bool


def test_output_type_serialization_nested():
    # typing
    assert serialize_type(List[Union[str, int]]) == "typing.List[typing.Union[str, int]]"
    assert serialize_type(List[Optional[str]]) == "typing.List[typing.Optional[str]]"
    assert serialize_type(List[Dict[str, int]]) == "typing.List[typing.Dict[str, int]]"
    assert serialize_type(typing.List[Dict[str, int]]) == "typing.List[typing.Dict[str, int]]"
    # builtins
    assert serialize_type(list[Union[str, int]]) == "list[typing.Union[str, int]]"
    assert serialize_type(list[Optional[str]]) == "list[typing.Optional[str]]"
    assert serialize_type(list[dict[str, int]]) == "list[dict[str, int]]"
    assert serialize_type(list[list[int]]) == "list[list[int]]"
    assert serialize_type(list[list[list[int]]]) == "list[list[list[int]]]"
    # PEP 604
    assert serialize_type(list[str | int]) == "list[str | int]"
    assert serialize_type(list[str | None]) == "list[str | None]"
    assert serialize_type(dict[str, int | None]) == "dict[str, int | None]"
    assert serialize_type(list[dict[str, int] | None]) == "list[dict[str, int] | None]"


def test_output_type_deserialization_nested():
    # typing
    assert deserialize_type("typing.List[typing.Union[str, int]]") == List[Union[str, int]]
    assert deserialize_type("typing.List[typing.Optional[str]]") == List[Optional[str]]
    assert deserialize_type("typing.List[typing.Dict[str, typing.List[int]]]") == List[Dict[str, List[int]]]
    assert deserialize_type("typing.List[typing.Dict[str, int]]") == typing.List[Dict[str, int]]
    # builtins
    assert deserialize_type("list[typing.Union[str, int]]") == list[Union[str, int]]
    assert deserialize_type("list[typing.Optional[str]]") == list[Optional[str]]
    assert deserialize_type("list[dict[str, list[int]]]") == list[dict[str, list[int]]]
    assert deserialize_type("list[dict[str, int]]") == list[dict[str, int]]
    assert deserialize_type("list[list[int]]") == list[list[int]]
    assert deserialize_type("list[list[list[int]]]") == list[list[list[int]]]
    # PEP 604
    assert deserialize_type("list[str | int]") == list[Union[str, int]]
    assert deserialize_type("list[str | None]") == list[Union[str, None]]
    assert deserialize_type("dict[str, int | None]") == dict[str, Union[int, None]]
    assert deserialize_type("list[dict[str, int] | None]") == list[Union[dict[str, int], None]]


def test_output_type_serialization_haystack_dataclasses():
    # typing
    # Answer
    assert serialize_type(Answer) == "haystack.dataclasses.answer.Answer"
    assert serialize_type(List[Answer]) == "typing.List[haystack.dataclasses.answer.Answer]"
    assert serialize_type(typing.Dict[int, Answer]) == "typing.Dict[int, haystack.dataclasses.answer.Answer]"
    # Bytestream
    assert serialize_type(ByteStream) == "haystack.dataclasses.byte_stream.ByteStream"
    assert serialize_type(List[ByteStream]) == "typing.List[haystack.dataclasses.byte_stream.ByteStream]"
    assert (
        serialize_type(typing.Dict[int, ByteStream]) == "typing.Dict[int, haystack.dataclasses.byte_stream.ByteStream]"
    )
    # Chat Message
    assert serialize_type(ChatMessage) == "haystack.dataclasses.chat_message.ChatMessage"
    assert serialize_type(List[ChatMessage]) == "typing.List[haystack.dataclasses.chat_message.ChatMessage]"
    assert (
        serialize_type(typing.Dict[int, ChatMessage])
        == "typing.Dict[int, haystack.dataclasses.chat_message.ChatMessage]"
    )
    # Document
    assert serialize_type(Document) == "haystack.dataclasses.document.Document"
    assert serialize_type(List[Document]) == "typing.List[haystack.dataclasses.document.Document]"
    assert serialize_type(typing.Dict[int, Document]) == "typing.Dict[int, haystack.dataclasses.document.Document]"
    # builtins
    # Answer
    assert serialize_type(list[Answer]) == "list[haystack.dataclasses.answer.Answer]"
    assert serialize_type(dict[int, Answer]) == "dict[int, haystack.dataclasses.answer.Answer]"
    # Bytestream
    assert serialize_type(list[ByteStream]) == "list[haystack.dataclasses.byte_stream.ByteStream]"
    assert serialize_type(dict[int, ByteStream]) == "dict[int, haystack.dataclasses.byte_stream.ByteStream]"
    # Chat Message
    assert serialize_type(list[ChatMessage]) == "list[haystack.dataclasses.chat_message.ChatMessage]"
    assert serialize_type(dict[int, ChatMessage]) == "dict[int, haystack.dataclasses.chat_message.ChatMessage]"
    # Document
    assert serialize_type(list[Document]) == "list[haystack.dataclasses.document.Document]"
    assert serialize_type(dict[int, Document]) == "dict[int, haystack.dataclasses.document.Document]"


def test_output_type_deserialization_haystack_dataclasses():
    # typing
    # Answer
    assert deserialize_type("haystack.dataclasses.answer.Answer") == Answer
    assert deserialize_type("typing.List[haystack.dataclasses.answer.Answer]") == List[Answer]
    assert deserialize_type("typing.Dict[int, haystack.dataclasses.answer.Answer]") == typing.Dict[int, Answer]
    # ByteStream
    assert deserialize_type("haystack.dataclasses.byte_stream.ByteStream") == ByteStream
    assert deserialize_type("typing.List[haystack.dataclasses.byte_stream.ByteStream]") == List[ByteStream]
    assert (
        deserialize_type("typing.Dict[int, haystack.dataclasses.byte_stream.ByteStream]")
        == typing.Dict[int, ByteStream]
    )
    # Chat Message
    assert deserialize_type("typing.List[haystack.dataclasses.chat_message.ChatMessage]") == typing.List[ChatMessage]
    assert (
        deserialize_type("typing.Dict[int, haystack.dataclasses.chat_message.ChatMessage]")
        == typing.Dict[int, ChatMessage]
    )
    assert deserialize_type("haystack.dataclasses.chat_message.ChatMessage") == ChatMessage
    # Document
    assert deserialize_type("haystack.dataclasses.document.Document") == Document
    assert deserialize_type("typing.List[haystack.dataclasses.document.Document]") == typing.List[Document]
    assert deserialize_type("typing.Dict[int, haystack.dataclasses.document.Document]") == typing.Dict[int, Document]
    # builtins
    # Answer
    assert deserialize_type("list[haystack.dataclasses.answer.Answer]") == list[Answer]
    assert deserialize_type("dict[int, haystack.dataclasses.answer.Answer]") == dict[int, Answer]
    # ByteStream
    assert deserialize_type("list[haystack.dataclasses.byte_stream.ByteStream]") == list[ByteStream]
    assert deserialize_type("dict[int, haystack.dataclasses.byte_stream.ByteStream]") == dict[int, ByteStream]
    # Chat Message
    assert deserialize_type("list[haystack.dataclasses.chat_message.ChatMessage]") == list[ChatMessage]
    assert deserialize_type("dict[int, haystack.dataclasses.chat_message.ChatMessage]") == dict[int, ChatMessage]
    # Document
    assert deserialize_type("list[haystack.dataclasses.document.Document]") == list[Document]
    assert deserialize_type("dict[int, haystack.dataclasses.document.Document]") == dict[int, Document]


def test_output_type_serialization_pep_604():
    # PEP 604 allows for union types to be defined with the `|` operator
    assert serialize_type(str | int) == "str | int"
    assert serialize_type(List[str] | List[int]) == "typing.Union[typing.List[str], typing.List[int]]"
    assert (
        serialize_type(Dict[str, int] | Dict[int, str]) == "typing.Union[typing.Dict[str, int], typing.Dict[int, str]]"
    )
    assert serialize_type(str | None) == "str | None"
    assert serialize_type(list[str] | None) == "list[str] | None"
    assert serialize_type(int | float | str) == "int | float | str"
    assert serialize_type(dict[str, int] | None) == "dict[str, int] | None"
    assert serialize_type(set[int] | None) == "set[int] | None"
    assert serialize_type(tuple[int, str] | None) == "tuple[int, str] | None"
    assert serialize_type(list[int] | list[str]) == "list[int] | list[str]"
    assert serialize_type(dict[str, int] | dict[int, str]) == "dict[str, int] | dict[int, str]"


def test_output_type_deserialization_pep_604():
    assert deserialize_type("str | int") == Union[str, int]
    assert deserialize_type("str | None") == Union[str, None]
    assert deserialize_type("int | float") == Union[int, float]
    assert deserialize_type("str | int | float") == Union[str, int, float]
    assert deserialize_type("str | int | None") == Union[str, int, None]
    assert deserialize_type("list[str] | None") == Union[list[str], None]
    assert deserialize_type("list[str] | list[int]") == Union[list[str], list[int]]
    assert deserialize_type("dict[str, int] | None") == Union[dict[str, int], None]
    assert deserialize_type("list[dict[str, int]] | None") == Union[list[dict[str, int]], None]
    assert deserialize_type("dict[str, list[int]] | set[str]") == Union[dict[str, list[int]], set[str]]
    assert deserialize_type("typing.List[str] | None") == Union[List[str], None]
    assert deserialize_type("typing.Dict[str, int] | typing.List[str]") == Union[Dict[str, int], List[str]]
    assert deserialize_type("set[int] | None") == Union[set[int], None]
    assert deserialize_type("tuple[int, str] | None") == Union[tuple[int, str], None]
    assert deserialize_type("frozenset[int] | None") == Union[frozenset[int], None]
    assert deserialize_type("dict[str, int] | dict[int, str]") == Union[dict[str, int], dict[int, str]]
    assert deserialize_type("list[int] | list[str] | list[float]") == Union[list[int], list[str], list[float]]


def test_is_union_type():
    assert _is_union_type(Union) is True
    assert _is_union_type(UnionType) is True
    assert _is_union_type(Union[str, int]) is True
    assert _is_union_type(str | int) is True
    assert _is_union_type(str | None) is True
    assert _is_union_type(Optional[str]) is True

    assert _is_union_type(str) is False
    assert _is_union_type(None) is False
    assert _is_union_type(list[str]) is False
    assert _is_union_type(dict[str, int]) is False


def test_parse_pep604_union_args():
    assert _parse_pep604_union_args("str | int") == ["str", "int"]
    assert _parse_pep604_union_args("str | None") == ["str", "None"]
    assert _parse_pep604_union_args("str | int | float") == ["str", "int", "float"]
    assert _parse_pep604_union_args("str | int | None") == ["str", "int", "None"]

    # Nested generics
    assert _parse_pep604_union_args("list[str] | None") == ["list[str]", "None"]
    assert _parse_pep604_union_args("list[str] | dict[str, int]") == ["list[str]", "dict[str, int]"]
    assert _parse_pep604_union_args("list[str] | dict[str, int] | None") == ["list[str]", "dict[str, int]", "None"]
    assert _parse_pep604_union_args("set[int] | None") == ["set[int]", "None"]
    assert _parse_pep604_union_args("tuple[int, str] | None") == ["tuple[int, str]", "None"]
    assert _parse_pep604_union_args("dict[str, list[int]] | set[str]") == ["dict[str, list[int]]", "set[str]"]
    assert _parse_pep604_union_args("list[int] | list[str] | list[float]") == ["list[int]", "list[str]", "list[float]"]


def test_build_pep604_union_type():
    result = _build_pep604_union_type([str])
    assert result == str

    result = _build_pep604_union_type([str, int])
    assert result == str | int

    result = _build_pep604_union_type([str, int, float])
    assert result == str | int | float

    result = _build_pep604_union_type([str, type(None)])
    assert result == str | None

    result = _build_pep604_union_type([list[str], dict[str, int]])
    assert result == list[str] | dict[str, int]
