# 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 operator
import os
from typing import TYPE_CHECKING
from unittest import mock
from unittest.mock import MagicMock

import pytest

from airflow.configuration import conf
from airflow.models.dagrun import DagRun, DagRunType
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.models.xcom import BaseXCom, XCom, resolve_xcom_backend
from airflow.operators.empty import EmptyOperator
from airflow.settings import json
from airflow.utils import timezone
from airflow.utils.session import create_session
from airflow.utils.xcom import XCOM_RETURN_KEY
from tests.test_utils.config import conf_vars

if TYPE_CHECKING:
    from sqlalchemy.orm import Session


class CustomXCom(BaseXCom):
    orm_deserialize_value = mock.Mock()


@pytest.fixture(autouse=True)
def reset_db():
    """Reset XCom entries."""
    with create_session() as session:
        session.query(DagRun).delete()
        session.query(XCom).delete()


@pytest.fixture()
def task_instance_factory(request, session: Session):
    def func(*, dag_id, task_id, execution_date):
        run_id = DagRun.generate_run_id(DagRunType.SCHEDULED, execution_date)
        run = DagRun(
            dag_id=dag_id,
            run_type=DagRunType.SCHEDULED,
            run_id=run_id,
            execution_date=execution_date,
        )
        session.add(run)
        ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id)
        ti.dag_id = dag_id
        session.add(ti)
        session.commit()

        def cleanup_database():
            # This should also clear task instances by cascading.
            session.query(DagRun).filter_by(id=run.id).delete()
            session.commit()

        request.addfinalizer(cleanup_database)
        return ti

    return func


@pytest.fixture()
def task_instance(task_instance_factory):
    return task_instance_factory(
        dag_id="dag",
        task_id="task_1",
        execution_date=timezone.datetime(2021, 12, 3, 4, 56),
    )


@pytest.fixture()
def task_instances(session, task_instance):
    ti2 = TaskInstance(EmptyOperator(task_id="task_2"), run_id=task_instance.run_id)
    ti2.dag_id = task_instance.dag_id
    session.add(ti2)
    session.commit()
    return task_instance, ti2  # ti2 will be cleaned up automatically with the DAG run.


class TestXCom:
    @conf_vars({("core", "xcom_backend"): "tests.models.test_xcom.CustomXCom"})
    def test_resolve_xcom_class(self):
        cls = resolve_xcom_backend()
        assert issubclass(cls, CustomXCom)

    @conf_vars({("core", "xcom_backend"): "", ("core", "enable_xcom_pickling"): "False"})
    def test_resolve_xcom_class_fallback_to_basexcom(self):
        cls = resolve_xcom_backend()
        assert issubclass(cls, BaseXCom)
        assert cls.serialize_value([1]) == b"[1]"

    @conf_vars({("core", "enable_xcom_pickling"): "False"})
    @conf_vars({("core", "xcom_backend"): "to be removed"})
    def test_resolve_xcom_class_fallback_to_basexcom_no_config(self):
        conf.remove_option("core", "xcom_backend")
        cls = resolve_xcom_backend()
        assert issubclass(cls, BaseXCom)
        assert cls.serialize_value([1]) == b"[1]"

    def test_xcom_deserialize_with_json_to_pickle_switch(self, task_instance, session):
        ti_key = TaskInstanceKey(
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
        )
        with conf_vars({("core", "enable_xcom_pickling"): "False"}):
            XCom.set(
                key="xcom_test3",
                value={"key": "value"},
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                run_id=task_instance.run_id,
                session=session,
            )
        with conf_vars({("core", "enable_xcom_pickling"): "True"}):
            ret_value = XCom.get_value(key="xcom_test3", ti_key=ti_key, session=session)
        assert ret_value == {"key": "value"}

    def test_xcom_deserialize_with_pickle_to_json_switch(self, task_instance, session):
        with conf_vars({("core", "enable_xcom_pickling"): "True"}):
            XCom.set(
                key="xcom_test3",
                value={"key": "value"},
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                run_id=task_instance.run_id,
                session=session,
            )
        with conf_vars({("core", "enable_xcom_pickling"): "False"}):
            ret_value = XCom.get_one(
                key="xcom_test3",
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                run_id=task_instance.run_id,
                session=session,
            )
        assert ret_value == {"key": "value"}

    @conf_vars({("core", "xcom_enable_pickling"): "False"})
    def test_xcom_disable_pickle_type_fail_on_non_json(self, task_instance, session):
        class PickleRce:
            def __reduce__(self):
                return os.system, ("ls -alt",)

        with pytest.raises(TypeError):
            XCom.set(
                key="xcom_test3",
                value=PickleRce(),
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                run_id=task_instance.run_id,
                session=session,
            )

    @mock.patch("airflow.models.xcom.XCom.orm_deserialize_value")
    def test_xcom_init_on_load_uses_orm_deserialize_value(self, mock_orm_deserialize):
        instance = BaseXCom(
            key="key",
            value="value",
            timestamp=timezone.utcnow(),
            execution_date=timezone.utcnow(),
            task_id="task_id",
            dag_id="dag_id",
        )
        instance.init_on_load()
        mock_orm_deserialize.assert_called_once_with()

    @conf_vars({("core", "xcom_backend"): "tests.models.test_xcom.CustomXCom"})
    def test_get_one_custom_backend_no_use_orm_deserialize_value(self, task_instance, session):
        """Test that XCom.get_one does not call orm_deserialize_value"""
        XCom = resolve_xcom_backend()
        XCom.set(
            key=XCOM_RETURN_KEY,
            value={"key": "value"},
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        )

        value = XCom.get_one(
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        )
        assert value == {"key": "value"}
        XCom.orm_deserialize_value.assert_not_called()

    @conf_vars({("core", "enable_xcom_pickling"): "False"})
    @mock.patch("airflow.models.xcom.conf.getimport")
    def test_set_serialize_call_old_signature(self, get_import, task_instance):
        """
        When XCom.serialize_value takes only param ``value``, other kwargs should be ignored.
        """
        serialize_watcher = MagicMock()

        class OldSignatureXCom(BaseXCom):
            @staticmethod
            def serialize_value(value, **kwargs):
                serialize_watcher(value=value, **kwargs)
                return json.dumps(value).encode("utf-8")

        get_import.return_value = OldSignatureXCom

        XCom = resolve_xcom_backend()
        XCom.set(
            key=XCOM_RETURN_KEY,
            value={"my_xcom_key": "my_xcom_value"},
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
        )
        serialize_watcher.assert_called_once_with(value={"my_xcom_key": "my_xcom_value"})

    @conf_vars({("core", "enable_xcom_pickling"): "False"})
    @mock.patch("airflow.models.xcom.conf.getimport")
    def test_set_serialize_call_current_signature(self, get_import, task_instance):
        """
        When XCom.serialize_value includes params execution_date, key, dag_id, task_id and run_id,
        then XCom.set should pass all of them.
        """
        serialize_watcher = MagicMock()

        class CurrentSignatureXCom(BaseXCom):
            @staticmethod
            def serialize_value(
                value,
                key=None,
                dag_id=None,
                task_id=None,
                run_id=None,
                map_index=None,
            ):
                serialize_watcher(
                    value=value,
                    key=key,
                    dag_id=dag_id,
                    task_id=task_id,
                    run_id=run_id,
                    map_index=map_index,
                )
                return json.dumps(value).encode("utf-8")

        get_import.return_value = CurrentSignatureXCom

        XCom = resolve_xcom_backend()
        XCom.set(
            key=XCOM_RETURN_KEY,
            value={"my_xcom_key": "my_xcom_value"},
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            map_index=-1,
        )
        serialize_watcher.assert_called_once_with(
            key=XCOM_RETURN_KEY,
            value={"my_xcom_key": "my_xcom_value"},
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            map_index=-1,
        )


@pytest.fixture(
    params=[
        pytest.param("true", id="enable_xcom_pickling=true"),
        pytest.param("false", id="enable_xcom_pickling=false"),
    ],
)
def setup_xcom_pickling(request):
    with conf_vars({("core", "enable_xcom_pickling"): str(request.param)}):
        yield


@pytest.fixture()
def push_simple_json_xcom(session):
    def func(*, ti: TaskInstance, key: str, value):
        return XCom.set(
            key=key,
            value=value,
            dag_id=ti.dag_id,
            task_id=ti.task_id,
            run_id=ti.run_id,
            session=session,
        )

    return func


@pytest.mark.usefixtures("setup_xcom_pickling")
class TestXComGet:
    @pytest.fixture()
    def setup_for_xcom_get_one(self, task_instance, push_simple_json_xcom):
        push_simple_json_xcom(ti=task_instance, key="xcom_1", value={"key": "value"})

    @pytest.mark.usefixtures("setup_for_xcom_get_one")
    def test_xcom_get_one(self, session, task_instance):
        stored_value = XCom.get_one(
            key="xcom_1",
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        )
        assert stored_value == {"key": "value"}

    @pytest.mark.usefixtures("setup_for_xcom_get_one")
    def test_xcom_get_one_with_execution_date(self, session, task_instance):
        with pytest.deprecated_call():
            stored_value = XCom.get_one(
                key="xcom_1",
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                execution_date=task_instance.execution_date,
                session=session,
            )
        assert stored_value == {"key": "value"}

    @pytest.fixture()
    def tis_for_xcom_get_one_from_prior_date(self, task_instance_factory, push_simple_json_xcom):
        date1 = timezone.datetime(2021, 12, 3, 4, 56)
        ti1 = task_instance_factory(dag_id="dag", execution_date=date1, task_id="task_1")
        ti2 = task_instance_factory(
            dag_id="dag",
            execution_date=date1 + datetime.timedelta(days=1),
            task_id="task_1",
        )

        # The earlier run pushes an XCom, but not the later run, but the later
        # run can get this earlier XCom with ``include_prior_dates``.
        push_simple_json_xcom(ti=ti1, key="xcom_1", value={"key": "value"})

        return ti1, ti2

    def test_xcom_get_one_from_prior_date(self, session, tis_for_xcom_get_one_from_prior_date):
        _, ti2 = tis_for_xcom_get_one_from_prior_date
        retrieved_value = XCom.get_one(
            run_id=ti2.run_id,
            key="xcom_1",
            task_id="task_1",
            dag_id="dag",
            include_prior_dates=True,
            session=session,
        )
        assert retrieved_value == {"key": "value"}

    def test_xcom_get_one_from_prior_with_execution_date(
        self,
        session,
        tis_for_xcom_get_one_from_prior_date,
    ):
        _, ti2 = tis_for_xcom_get_one_from_prior_date
        with pytest.deprecated_call():
            retrieved_value = XCom.get_one(
                execution_date=ti2.execution_date,
                key="xcom_1",
                task_id="task_1",
                dag_id="dag",
                include_prior_dates=True,
                session=session,
            )
        assert retrieved_value == {"key": "value"}

    @pytest.fixture()
    def setup_for_xcom_get_many_single_argument_value(self, task_instance, push_simple_json_xcom):
        push_simple_json_xcom(ti=task_instance, key="xcom_1", value={"key": "value"})

    @pytest.mark.usefixtures("setup_for_xcom_get_many_single_argument_value")
    def test_xcom_get_many_single_argument_value(self, session, task_instance):
        stored_xcoms = XCom.get_many(
            key="xcom_1",
            dag_ids=task_instance.dag_id,
            task_ids=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        ).all()
        assert len(stored_xcoms) == 1
        assert stored_xcoms[0].key == "xcom_1"
        assert stored_xcoms[0].value == {"key": "value"}

    @pytest.mark.usefixtures("setup_for_xcom_get_many_single_argument_value")
    def test_xcom_get_many_single_argument_value_with_execution_date(self, session, task_instance):
        with pytest.deprecated_call():
            stored_xcoms = XCom.get_many(
                execution_date=task_instance.execution_date,
                key="xcom_1",
                dag_ids=task_instance.dag_id,
                task_ids=task_instance.task_id,
                session=session,
            ).all()
        assert len(stored_xcoms) == 1
        assert stored_xcoms[0].key == "xcom_1"
        assert stored_xcoms[0].value == {"key": "value"}

    @pytest.fixture()
    def setup_for_xcom_get_many_multiple_tasks(self, task_instances, push_simple_json_xcom):
        ti1, ti2 = task_instances
        push_simple_json_xcom(ti=ti1, key="xcom_1", value={"key1": "value1"})
        push_simple_json_xcom(ti=ti2, key="xcom_1", value={"key2": "value2"})

    @pytest.mark.usefixtures("setup_for_xcom_get_many_multiple_tasks")
    def test_xcom_get_many_multiple_tasks(self, session, task_instance):
        stored_xcoms = XCom.get_many(
            key="xcom_1",
            dag_ids=task_instance.dag_id,
            task_ids=["task_1", "task_2"],
            run_id=task_instance.run_id,
            session=session,
        )
        sorted_values = [x.value for x in sorted(stored_xcoms, key=operator.attrgetter("task_id"))]
        assert sorted_values == [{"key1": "value1"}, {"key2": "value2"}]

    @pytest.mark.usefixtures("setup_for_xcom_get_many_multiple_tasks")
    def test_xcom_get_many_multiple_tasks_with_execution_date(self, session, task_instance):
        with pytest.deprecated_call():
            stored_xcoms = XCom.get_many(
                execution_date=task_instance.execution_date,
                key="xcom_1",
                dag_ids=task_instance.dag_id,
                task_ids=["task_1", "task_2"],
                session=session,
            )
        sorted_values = [x.value for x in sorted(stored_xcoms, key=operator.attrgetter("task_id"))]
        assert sorted_values == [{"key1": "value1"}, {"key2": "value2"}]

    @pytest.fixture()
    def tis_for_xcom_get_many_from_prior_dates(self, task_instance_factory, push_simple_json_xcom):
        date1 = timezone.datetime(2021, 12, 3, 4, 56)
        date2 = date1 + datetime.timedelta(days=1)
        ti1 = task_instance_factory(dag_id="dag", task_id="task_1", execution_date=date1)
        ti2 = task_instance_factory(dag_id="dag", task_id="task_1", execution_date=date2)
        push_simple_json_xcom(ti=ti1, key="xcom_1", value={"key1": "value1"})
        push_simple_json_xcom(ti=ti2, key="xcom_1", value={"key2": "value2"})
        return ti1, ti2

    def test_xcom_get_many_from_prior_dates(self, session, tis_for_xcom_get_many_from_prior_dates):
        ti1, ti2 = tis_for_xcom_get_many_from_prior_dates
        stored_xcoms = XCom.get_many(
            run_id=ti2.run_id,
            key="xcom_1",
            dag_ids="dag",
            task_ids="task_1",
            include_prior_dates=True,
            session=session,
        )

        # The retrieved XComs should be ordered by logical date, latest first.
        assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}]
        assert [x.execution_date for x in stored_xcoms] == [ti2.execution_date, ti1.execution_date]

    def test_xcom_get_many_from_prior_dates_with_execution_date(
        self,
        session,
        tis_for_xcom_get_many_from_prior_dates,
    ):
        ti1, ti2 = tis_for_xcom_get_many_from_prior_dates
        with pytest.deprecated_call():
            stored_xcoms = XCom.get_many(
                execution_date=ti2.execution_date,
                key="xcom_1",
                dag_ids="dag",
                task_ids="task_1",
                include_prior_dates=True,
                session=session,
            )

        # The retrieved XComs should be ordered by logical date, latest first.
        assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}]
        assert [x.execution_date for x in stored_xcoms] == [ti2.execution_date, ti1.execution_date]


@pytest.mark.usefixtures("setup_xcom_pickling")
class TestXComSet:
    def test_xcom_set(self, session, task_instance):
        XCom.set(
            key="xcom_1",
            value={"key": "value"},
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        )
        stored_xcoms = session.query(XCom).all()
        assert stored_xcoms[0].key == "xcom_1"
        assert stored_xcoms[0].value == {"key": "value"}
        assert stored_xcoms[0].dag_id == "dag"
        assert stored_xcoms[0].task_id == "task_1"
        assert stored_xcoms[0].execution_date == task_instance.execution_date

    def test_xcom_set_with_execution_date(self, session, task_instance):
        with pytest.deprecated_call():
            XCom.set(
                key="xcom_1",
                value={"key": "value"},
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                execution_date=task_instance.execution_date,
                session=session,
            )
        stored_xcoms = session.query(XCom).all()
        assert stored_xcoms[0].key == "xcom_1"
        assert stored_xcoms[0].value == {"key": "value"}
        assert stored_xcoms[0].dag_id == "dag"
        assert stored_xcoms[0].task_id == "task_1"
        assert stored_xcoms[0].execution_date == task_instance.execution_date

    @pytest.fixture()
    def setup_for_xcom_set_again_replace(self, task_instance, push_simple_json_xcom):
        push_simple_json_xcom(ti=task_instance, key="xcom_1", value={"key1": "value1"})

    @pytest.mark.usefixtures("setup_for_xcom_set_again_replace")
    def test_xcom_set_again_replace(self, session, task_instance):
        assert session.query(XCom).one().value == {"key1": "value1"}
        XCom.set(
            key="xcom_1",
            value={"key2": "value2"},
            dag_id=task_instance.dag_id,
            task_id="task_1",
            run_id=task_instance.run_id,
            session=session,
        )
        assert session.query(XCom).one().value == {"key2": "value2"}

    @pytest.mark.usefixtures("setup_for_xcom_set_again_replace")
    def test_xcom_set_again_replace_with_execution_date(self, session, task_instance):
        assert session.query(XCom).one().value == {"key1": "value1"}
        with pytest.deprecated_call():
            XCom.set(
                key="xcom_1",
                value={"key2": "value2"},
                dag_id=task_instance.dag_id,
                task_id="task_1",
                execution_date=task_instance.execution_date,
                session=session,
            )
        assert session.query(XCom).one().value == {"key2": "value2"}


@pytest.mark.usefixtures("setup_xcom_pickling")
class TestXComClear:
    @pytest.fixture()
    def setup_for_xcom_clear(self, task_instance, push_simple_json_xcom):
        push_simple_json_xcom(ti=task_instance, key="xcom_1", value={"key": "value"})

    @pytest.mark.usefixtures("setup_for_xcom_clear")
    def test_xcom_clear(self, session, task_instance):
        assert session.query(XCom).count() == 1
        XCom.clear(
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id=task_instance.run_id,
            session=session,
        )
        assert session.query(XCom).count() == 0

    @pytest.mark.usefixtures("setup_for_xcom_clear")
    def test_xcom_clear_with_execution_date(self, session, task_instance):
        assert session.query(XCom).count() == 1
        with pytest.deprecated_call():
            XCom.clear(
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                execution_date=task_instance.execution_date,
                session=session,
            )
        assert session.query(XCom).count() == 0

    @pytest.mark.usefixtures("setup_for_xcom_clear")
    def test_xcom_clear_different_run(self, session, task_instance):
        XCom.clear(
            dag_id=task_instance.dag_id,
            task_id=task_instance.task_id,
            run_id="different_run",
            session=session,
        )
        assert session.query(XCom).count() == 1

    @pytest.mark.usefixtures("setup_for_xcom_clear")
    def test_xcom_clear_different_execution_date(self, session, task_instance):
        with pytest.deprecated_call():
            XCom.clear(
                dag_id=task_instance.dag_id,
                task_id=task_instance.task_id,
                execution_date=timezone.utcnow(),
                session=session,
            )
        assert session.query(XCom).count() == 1
