# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import datetime
import enum
import textwrap
from collections import namedtuple
from dataclasses import dataclass
from importlib import import_module, metadata
from typing import ClassVar

import attr
import pytest
from packaging import version
from pydantic import BaseModel, create_model

from airflow._shared.module_loading import import_string, iter_namespace, qualname
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.serde import (
    CLASSNAME,
    DATA,
    SCHEMA_ID,
    VERSION,
    _extra_allowed,
    _get_patterns,
    _get_regexp_patterns,
    _match,
    _match_glob,
    _match_regexp,
    allow_class,
    deserialize,
    iter_pydantic_models,
    serialize,
)

from tests_common.test_utils.config import conf_vars


@pytest.fixture
def recalculate_patterns():
    _get_patterns.cache_clear()
    _get_regexp_patterns.cache_clear()
    _match_glob.cache_clear()
    _match_regexp.cache_clear()
    try:
        yield
    finally:
        _get_patterns.cache_clear()
        _get_regexp_patterns.cache_clear()
        _match_glob.cache_clear()
        _match_regexp.cache_clear()


def generate_serializers_importable_tests():
    """
    Generate test cases for `test_serializers_importable_and_str`.

    The function iterates through all the modules defined under `airflow.sdk.serde.serializers`. It loads
    the import strings defined in the `serializers` from each module, and create a test case to verify that the
    serializer is importable.
    """
    import airflow.sdk.serde.serializers

    NUMPY_VERSION = version.parse(metadata.version("numpy"))

    serializer_tests = []

    for _, name, _ in iter_namespace(airflow.sdk.serde.serializers):
        ############################################################
        # Handle compatibility / optional dependency at module level
        ############################################################
        # https://github.com/apache/airflow/pull/37320
        if name == "airflow.sdk.serde.serializers.iceberg":
            try:
                import pyiceberg  # noqa: F401
            except ImportError:
                continue
        # https://github.com/apache/airflow/pull/38074
        if name == "airflow.sdk.serde.serializers.deltalake":
            try:
                import deltalake  # noqa: F401
            except ImportError:
                continue
        mod = import_module(name)
        for s in getattr(mod, "serializers", list()):
            ############################################################
            # Handle compatibility issue at serializer level
            ############################################################
            if s == "numpy.bool" and NUMPY_VERSION.major < 2:
                reason = textwrap.dedent(f"""\
                    Current NumPy version: {NUMPY_VERSION}

                    In NumPy 1.20, `numpy.bool` was deprecated as an alias for the built-in `bool`.
                    For NumPy versions <= 1.26, attempting to import `numpy.bool` raises an ImportError.
                    Starting with NumPy 2.0, `numpy.bool` is reintroduced as the NumPy scalar type,
                    and `numpy.bool_` becomes an alias for `numpy.bool`.

                    The serializers are loaded lazily at runtime. As a result:
                    - With NumPy <= 1.26, only `numpy.bool_` is loaded.
                    - With NumPy >= 2.0, only `numpy.bool` is loaded.

                    This test case deliberately attempts to import both `numpy.bool` and `numpy.bool_`,
                    regardless of the installed NumPy version. Therefore, when NumPy <= 1.26 is installed,
                    importing `numpy.bool` will raise an ImportError.
                """)
                serializer_tests.append(pytest.param(name, s, marks=pytest.mark.skip(reason=reason)))
            else:
                serializer_tests.append(pytest.param(name, s))
    return serializer_tests


SERIALIZER_TESTS = generate_serializers_importable_tests()


class Z:
    __version__: ClassVar[int] = 1

    def __init__(self, x):
        self.x = x

    def serialize(self) -> dict:
        return dict({"x": self.x})

    @staticmethod
    def deserialize(data: dict, version: int):
        if version != 1:
            raise TypeError("version != 1")
        return Z(data["x"])

    def __eq__(self, other):
        return self.x == other.x

    def __hash__(self):
        return hash(self.x)


@attr.define
class Y:
    x: int
    __version__: ClassVar[int] = 1

    def __init__(self, x):
        self.x = x


class X:
    pass


@dataclass
class W:
    __version__: ClassVar[int] = 2
    x: int


@dataclass
class V:
    __version__: ClassVar[int] = 1
    w: W
    s: list
    t: tuple
    c: int


class U(BaseModel):
    __version__: ClassVar[int] = 1
    x: int
    v: V
    u: tuple


@attr.define
class T:
    x: int
    y: Y
    u: tuple
    w: W

    __version__: ClassVar[int] = 1


class C:
    def __call__(self):
        return None


class NestedLeaf(BaseModel):
    value: int


class NestedMid(BaseModel):
    leaves: list[NestedLeaf]


class NestedRoot(BaseModel):
    mid: NestedMid
    maybe: NestedLeaf | None


class SelfRefNode(BaseModel):
    name: str
    children: list[SelfRefNode] = []


SelfRefNode.model_rebuild()


@pytest.mark.usefixtures("recalculate_patterns")
class TestSerDe:
    def test_ser_primitives(self):
        i = 10
        e = serialize(i)
        assert i == e

        i = 10.1
        e = serialize(i)
        assert i == e

        i = "test"
        e = serialize(i)
        assert i == e

        i = True
        e = serialize(i)
        assert i == e

        Color = enum.IntEnum("Color", ["RED", "GREEN"])
        i = Color.RED
        e = serialize(i)
        assert i == e

    def test_ser_collections(self):
        i = [1, 2]
        e = deserialize(serialize(i))
        assert i == e

        i = ("a", "b", "a", "c")
        e = deserialize(serialize(i))
        assert i == e

        i = {2, 3}
        e = deserialize(serialize(i))
        assert i == e

        i = frozenset({6, 7})
        e = deserialize(serialize(i))
        assert i == e

    def test_der_collections_compat(self):
        i = [1, 2]
        e = deserialize(i)
        assert i == e

        i = ("a", "b", "a", "c")
        e = deserialize(i)
        assert i == e

        i = {2, 3}
        e = deserialize(i)
        assert i == e

    def test_ser_plain_dict(self):
        i = {"a": 1, "b": 2}
        e = serialize(i)
        assert i == e

        i = {CLASSNAME: "cannot"}
        with pytest.raises(AttributeError, match="^reserved"):
            serialize(i)

        i = {SCHEMA_ID: "cannot"}
        with pytest.raises(AttributeError, match="^reserved"):
            serialize(i)

    def test_ser_namedtuple(self):
        CustomTuple = namedtuple("CustomTuple", ["id", "value"])
        data = CustomTuple(id=1, value="something")

        i = deserialize(serialize(data))
        e = (1, "something")
        assert i == e

    def test_no_serializer(self):
        i = Exception
        with pytest.raises(TypeError, match="^Cannot serialize"):
            serialize(i)

    def test_ser_registered(self):
        i = datetime.datetime(2000, 10, 1)
        e = serialize(i)
        assert e[DATA]

    def test_serder_custom(self):
        i = Z(1)
        e = serialize(i)
        assert Z.__version__ == e[VERSION]
        assert qualname(Z) == e[CLASSNAME]
        assert e[DATA]

        d = deserialize(e)
        assert i.x == getattr(d, "x", None)

    def test_serder_attr(self):
        i = Y(10)
        e = serialize(i)
        assert Y.__version__ == e[VERSION]
        assert qualname(Y) == e[CLASSNAME]
        assert e[DATA]

        d = deserialize(e)
        assert i.x == getattr(d, "x", None)

    def test_serder_dataclass(self):
        i = W(12)
        e = serialize(i)
        assert W.__version__ == e[VERSION]
        assert qualname(W) == e[CLASSNAME]
        assert e[DATA]

        d = deserialize(e)
        assert i.x == getattr(d, "x", None)

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "airflow.*",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_for_imports(self):
        i = Z(10)
        e = serialize(i)
        with pytest.raises(ImportError) as ex:
            deserialize(e)

        assert f"{qualname(Z)} was not found in allow list" in str(ex.value)

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "airflow.*",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_for_deserialize_pydantic_model(self):
        # for Pydantic model to be deserialized, it must be in `allowed_deserialization_classes`
        i = U(x=7, v=V(w=W(x=42), s=["hello", "world"], t=(1, 2, 3), c=99), u=("extra", 123))
        e = serialize(i)
        with pytest.raises(ImportError) as ex:
            deserialize(e)

        assert f"{qualname(U)} was not found in allow list" in str(ex.value)

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "unit.airflow.*",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_match(self):
        assert _match("unit.airflow.deep")
        assert _match("unit.wrongpath") is False

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "unit.airflow.deep",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_match_class(self):
        """
        Test the match function when passing a full classname as
        allowed_deserialization_classes
        """
        assert _match("unit.airflow.deep")
        assert _match("unit.airflow.FALSE") is False

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "",
            ("core", "allowed_deserialization_classes_regexp"): r"unit\.airflow\..*",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_match_regexp(self):
        """
        Test the match function when passing a path as
        allowed_deserialization_classes_regexp with no glob pattern defined
        """
        assert _match("unit.airflow.deep")
        assert _match("unit.wrongpath") is False

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "",
            ("core", "allowed_deserialization_classes_regexp"): r"unit\.airflow\.deep",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_match_class_regexp(self):
        """
        Test the match function when passing a full classname as
        allowed_deserialization_classes_regexp with no glob pattern defined
        """
        assert _match("unit.airflow.deep")
        assert _match("unit.airflow.FALSE") is False

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "",
            ("core", "allowed_deserialization_classes_regexp"): r"unit\.airflow\.Variable",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_list_regexp_does_not_prefix_match(self):
        """
        A pattern without an explicit end anchor must not admit classes that share
        the pattern as a prefix. ``re.match`` would let ``unit.airflow.Variable_Malicious``
        through because it only anchors at the start of the string; ``re.fullmatch``
        rejects it. Patterns with ``.*`` at the end retain prefix-style behaviour.
        """
        assert _match("unit.airflow.Variable")
        assert _match("unit.airflow.Variable_Malicious") is False
        assert _match("unit.airflow.VariableSubclass") is False

    @conf_vars(
        {
            ("core", "allowed_deserialization_classes"): "airflow.*",
        }
    )
    @pytest.mark.usefixtures("recalculate_patterns")
    def test_allow_class_round_trips_pydantic_subclass(self):
        """``allow_class`` lets a Pydantic subclass round-trip without editing the allow-list config."""
        instance = U(x=7, v=V(w=W(x=42), s=["a", "b"], t=(1, 2), c=99), u=("z", 0))
        snapshot = set(_extra_allowed)
        try:
            assert qualname(U) not in _extra_allowed
            allow_class(U)
            assert qualname(U) in _extra_allowed

            restored = deserialize(serialize(instance))
            assert isinstance(restored, U)
            assert restored == instance
        finally:
            _extra_allowed.clear()
            _extra_allowed.update(snapshot)

    def test_allow_class_rejects_locals_qualname(self):
        """Nested-in-function classes have ``<locals>`` in qualname and cannot round-trip."""

        def _make():
            class Local(BaseModel):
                v: int

            return Local

        with pytest.raises(ValueError, match="defined inside a function"):
            allow_class(_make())

    def test_allow_class_rejects_class_with_mismatched_module_attr(self):
        """A class whose qualname does not import back to itself must be rejected."""
        Mismatched = create_model("DifferentName", x=(int, ...))

        with pytest.raises(ValueError, match="cannot be re-imported|does not resolve"):
            allow_class(Mismatched)

    def test_iter_pydantic_models_shapes(self):
        """iter_pydantic_models finds models in bare, optional/union, and container annotations."""

        class M1(BaseModel):
            a: int

        class M2(BaseModel):
            b: int

        assert set(iter_pydantic_models(M1)) == {M1}
        assert set(iter_pydantic_models(M1 | None)) == {M1}
        assert set(iter_pydantic_models(M1 | M2)) == {M1, M2}
        assert set(iter_pydantic_models(list[M1])) == {M1}
        assert set(iter_pydantic_models(dict[str, M2])) == {M2}
        assert set(iter_pydantic_models(list[M1 | M2 | None])) == {M1, M2}
        # Non-model annotations yield nothing.
        assert set(iter_pydantic_models(str)) == set()
        assert set(iter_pydantic_models(list[int])) == set()

    def test_iter_pydantic_models_recurses_into_fields(self):
        """Models nested inside a model's fields are yielded, not just the top-level type."""
        assert set(iter_pydantic_models(NestedRoot)) == {NestedRoot, NestedMid, NestedLeaf}
        # Holds when the declared type is a container of the root model.
        assert set(iter_pydantic_models(list[NestedRoot])) == {NestedRoot, NestedMid, NestedLeaf}

    def test_iter_pydantic_models_terminates_on_self_reference(self):
        """A self-referential model must not loop forever."""
        assert set(iter_pydantic_models(SelfRefNode)) == {SelfRefNode}

    def test_incompatible_version(self):
        data = dict(
            {
                "__classname__": Y.__module__ + "." + Y.__qualname__,
                "__version__": 2,
            }
        )
        with pytest.raises(TypeError, match="newer than"):
            deserialize(data)

    def test_raise_undeserializable(self):
        data = dict(
            {
                "__classname__": X.__module__ + "." + X.__qualname__,
                "__version__": 0,
            }
        )
        with pytest.raises(TypeError, match="No deserializer") as exc_info:
            deserialize(data)
        assert (
            "It must provide a `deserialize(data, version)` staticmethod (matching how it "
            "serializes), or be decorated with @dataclass or @attr.define." in str(exc_info.value)
        )
        assert "authoring-and-scheduling/serializers.html" in str(exc_info.value)

    def test_backwards_compat(self):
        """
        Verify deserialization of old-style encoded Xcom values including nested ones
        """
        uri = "s3://does/not/exist"
        data = {
            "__type": "airflow.sdk.definitions.asset.Asset",
            "__source": None,
            "__var": {
                "__var": {
                    "uri": uri,
                    "extra": {
                        "__var": {"hi": "bye"},
                        "__type": "dict",
                    },
                },
                "__type": "dict",
            },
        }
        asset = deserialize(data)
        assert asset.extra == {"hi": "bye"}
        assert asset.uri == uri

    def test_backwards_compat_wrapped(self):
        """
        Verify deserialization of old-style wrapped XCom value
        """
        i = {
            "extra": {"__var": {"hi": "bye"}, "__type": "dict"},
        }
        e = deserialize(i)
        assert e["extra"] == {"hi": "bye"}

    @pytest.mark.parametrize(
        ("old_type", "expected"),
        [
            ("tuple", (1, 2, 3)),
            ("set", {1, 2, 3}),
            ("frozenset", frozenset([1, 2, 3])),
        ],
    )
    def test_backwards_compat_builtin_collections(self, old_type, expected):
        """Verify deserialization of old-style builtin collections (tuple/set/frozenset)."""
        data = {"__type": old_type, "__var": [1, 2, 3]}
        result = deserialize(data)
        assert result == expected
        assert type(result) is type(expected)

    def test_backwards_compat_builtin_collection_nested(self):
        """Verify deserialization of old-style tuple nested inside a dict."""
        data = {
            "arg1": "hello",
            "arg2": {"__type": "tuple", "__var": [1, 2]},
        }
        result = deserialize(data)
        assert result == {"arg1": "hello", "arg2": (1, 2)}

    def test_backwards_compat_timedelta(self):
        """Verify deserialization of old-style timedelta."""
        import datetime

        data = {"__type": "timedelta", "__var": 3600.0}
        result = deserialize(data)
        assert result == datetime.timedelta(seconds=3600)

    def test_encode_asset(self):
        asset = Asset(uri="mytest://asset", name="test")
        obj = deserialize(serialize(asset))
        assert asset.uri == obj.uri

    @pytest.mark.parametrize(("name", "s"), SERIALIZER_TESTS)
    def test_serializers_importable_and_str(self, name, s):
        """Test if all distributed serializers are lazy loading and can be imported"""
        if not isinstance(s, str):
            raise TypeError(f"{s} is not of type str. This is required for lazy loading")
        try:
            import_string(s)
        except ImportError:
            raise AttributeError(f"{s} cannot be imported (located in {name})")

    @pytest.mark.parametrize(
        ("obj", "expected"),
        [
            (
                Z(10),
                {
                    "__classname__": "tests.task_sdk.serde.test_serde.Z",
                    "__version__": 1,
                    "__data__": {"x": 10},
                },
            ),
            (
                W(2),
                {
                    "__classname__": "tests.task_sdk.serde.test_serde.W",
                    "__version__": 2,
                    "__data__": {"x": 2},
                },
            ),
        ],
    )
    def test_serialized_data(self, obj, expected):
        assert expected == serialize(obj)

    def test_deserialize_non_serialized_data(self):
        i = Z(10)
        e = deserialize(i)
        assert i == e

    def test_attr(self):
        i = T(y=Y(10), u=(1, 2), x=10, w=W(11))
        e = serialize(i)
        s = deserialize(e)
        assert i == s

    def test_error_when_serializing_callable_without_name(self):
        i = C()
        with pytest.raises(
            TypeError,
            match="Cannot serialize object of type <class 'tests.task_sdk.serde.test_serde.C'>",
        ) as exc_info:
            serialize(i)
        assert (
            "Give it a `serialize()` method and a `deserialize(data, version)` staticmethod, or "
            "decorate the class with @dataclass or @attr.define." in str(exc_info.value)
        )
        assert "authoring-and-scheduling/serializers.html" in str(exc_info.value)