#
# 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 asyncio
import datetime
import importlib
import time
from threading import Thread
from unittest.mock import MagicMock, patch

import pendulum
import pytest

from airflow import DAG
from airflow.config_templates import airflow_local_settings
from airflow.jobs.job import Job
from airflow.jobs.triggerer_job_runner import TriggererJobRunner, TriggerRunner, setup_queue_listener
from airflow.logging_config import configure_logging
from airflow.models import DagModel, DagRun, TaskInstance, Trigger
from airflow.models.baseoperator import BaseOperator
from airflow.operators.empty import EmptyOperator
from airflow.operators.python import PythonOperator
from airflow.triggers.base import TriggerEvent
from airflow.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
from airflow.triggers.testing import FailureTrigger, SuccessTrigger
from airflow.utils import timezone
from airflow.utils.log.logging_mixin import RedirectStdHandler
from airflow.utils.log.trigger_handler import LocalQueueHandler
from airflow.utils.session import create_session
from airflow.utils.state import State, TaskInstanceState
from airflow.utils.types import DagRunType
from tests.core.test_logging_config import reset_logging
from tests.test_utils.db import clear_db_dags, clear_db_runs


class TimeDeltaTrigger_(TimeDeltaTrigger):
    def __init__(self, delta, filename):
        super().__init__(delta=delta)
        self.filename = filename
        self.delta = delta

    async def run(self):
        with open(self.filename, "a") as f:
            f.write("hi\n")
        async for event in super().run():
            yield event

    def serialize(self):
        return (
            "tests.jobs.test_triggerer_job.TimeDeltaTrigger_",
            {"delta": self.delta, "filename": self.filename},
        )


@pytest.fixture(autouse=True)
def clean_database():
    """Fixture that cleans the database before and after every test."""
    clear_db_runs()
    clear_db_dags()
    yield  # Test runs here
    clear_db_dags()
    clear_db_runs()


@pytest.fixture
def session():
    """Fixture that provides a SQLAlchemy session"""
    with create_session() as session:
        yield session


def create_trigger_in_db(session, trigger, operator=None):
    dag_model = DagModel(dag_id="test_dag")
    dag = DAG(dag_id=dag_model.dag_id, start_date=pendulum.datetime(2023, 1, 1))
    run = DagRun(
        dag_id=dag_model.dag_id,
        run_id="test_run",
        execution_date=pendulum.datetime(2023, 1, 1),
        run_type=DagRunType.MANUAL,
    )
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    if operator:
        operator.dag = dag
    else:
        operator = BaseOperator(task_id="test_ti", dag=dag)
    task_instance = TaskInstance(operator, execution_date=run.execution_date, run_id=run.run_id)
    task_instance.trigger_id = trigger_orm.id
    session.add(dag_model)
    session.add(run)
    session.add(trigger_orm)
    session.add(task_instance)
    session.commit()
    return dag_model, run, trigger_orm, task_instance


def test_trigger_logging_sensitive_info(session, capsys):
    """
    Checks that when a trigger fires, it doesn't log any sensitive
    information from arguments
    """

    class SensitiveArgOperator(BaseOperator):
        def __init__(self, password, **kwargs):
            self.password = password
            super().__init__(**kwargs)

    # Use a trigger that will immediately succeed
    trigger = SuccessTrigger()
    op = SensitiveArgOperator(task_id="sensitive_arg_task", password="some_password")
    create_trigger_in_db(session, trigger, operator=op)
    triggerer_job = Job()
    triggerer_job_runner = TriggererJobRunner(triggerer_job)
    triggerer_job_runner.load_triggers()
    # Now, start TriggerRunner up (and set it as a daemon thread during tests)
    triggerer_job_runner.daemon = True
    triggerer_job_runner.trigger_runner.start()
    try:
        # Wait for up to 3 seconds for it to fire and appear in the event queue
        for _ in range(30):
            if triggerer_job_runner.trigger_runner.events:
                assert list(triggerer_job_runner.trigger_runner.events) == [(1, TriggerEvent(True))]
                break
            time.sleep(0.1)
        else:
            pytest.fail("TriggerRunner never sent the trigger event out")
    finally:
        # We always have to stop the runner
        triggerer_job_runner.trigger_runner.stop = True
    # Since we have now an in-memory process of forwarding the logs to stdout,
    # give it more time for the trigger event to write the log.
    time.sleep(0.5)
    stdout = capsys.readouterr().out

    assert "test_dag/test_run/sensitive_arg_task/-1/1 (ID 1) starting" in stdout
    assert "some_password" not in stdout


def test_is_alive():
    """Checks the heartbeat logic"""
    # Current time
    triggerer_job = Job(heartrate=10, state=State.RUNNING)
    assert triggerer_job.is_alive()

    # Slightly old, but still fresh
    triggerer_job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=20)
    assert triggerer_job.is_alive()

    # Old enough to fail
    triggerer_job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=31)
    assert not triggerer_job.is_alive()

    # Completed state should not be alive
    triggerer_job.state = State.SUCCESS
    triggerer_job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=10)
    assert not triggerer_job.is_alive(), "Completed jobs even with recent heartbeat should not be alive"


def test_is_needed(session):
    """Checks the triggerer-is-needed logic"""
    # No triggers, no need
    triggerer_job = Job(heartrate=10, state=State.RUNNING)
    triggerer_job_runner = TriggererJobRunner(triggerer_job)
    assert triggerer_job_runner.is_needed() is False
    # Add a trigger, it's needed
    trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    session.add(trigger_orm)
    session.commit()
    assert triggerer_job_runner.is_needed() is True


def test_capacity_decode():
    """
    Tests that TriggererJob correctly sets capacity to a valid value passed in as a CLI arg,
    handles invalid args, or sets it to a default value if no arg is passed.
    """
    # Positive cases
    variants = [
        42,
        None,
    ]
    for input_str in variants:
        job = Job()
        job_runner = TriggererJobRunner(job, capacity=input_str)
        assert job_runner.capacity == input_str or job_runner.capacity == 1000

    # Negative cases
    variants = [
        "NAN",
        0.5,
        -42,
        4 / 2,  # Resolves to a float, in addition to being just plain weird
    ]
    for input_str in variants:
        with pytest.raises(ValueError):
            job = Job()
            TriggererJobRunner(job=job, capacity=input_str)


def test_trigger_lifecycle(session):
    """
    Checks that the triggerer will correctly see a new Trigger in the database
    and send it to the trigger runner, and then delete it when it vanishes.
    """
    # Use a trigger that will not fire for the lifetime of the test
    # (we want to avoid it firing and deleting itself)
    trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
    dag_model, run, trigger_orm, task_instance = create_trigger_in_db(session, trigger)
    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()
    # Make sure it turned up in TriggerRunner's queue
    assert [x for x, y in job_runner.trigger_runner.to_create] == [1]
    # Now, start TriggerRunner up (and set it as a daemon thread during tests)
    job_runner.daemon = True
    job_runner.trigger_runner.start()
    try:
        # Wait for up to 3 seconds for it to appear in the TriggerRunner's storage
        for _ in range(30):
            if job_runner.trigger_runner.triggers:
                assert list(job_runner.trigger_runner.triggers.keys()) == [1]
                break
            time.sleep(0.1)
        else:
            pytest.fail("TriggerRunner never created trigger")
        # OK, now remove it from the DB
        session.delete(trigger_orm)
        session.commit()
        # Re-load the triggers
        job_runner.load_triggers()
        # Wait for up to 3 seconds for it to vanish from the TriggerRunner's storage
        for _ in range(30):
            if not job_runner.trigger_runner.triggers:
                break
            time.sleep(0.1)
        else:
            pytest.fail("TriggerRunner never deleted trigger")
    finally:
        # We always have to stop the runner
        job_runner.trigger_runner.stop = True


class TestTriggerRunner:
    @pytest.mark.asyncio
    @patch("airflow.jobs.triggerer_job_runner.TriggerRunner.set_individual_trigger_logging")
    async def test_run_trigger_canceled(self, session) -> None:
        trigger_runner = TriggerRunner()
        trigger_runner.triggers = {1: {"task": MagicMock(), "name": "mock_name", "events": 0}}
        mock_trigger = MagicMock()
        mock_trigger.task_instance.trigger_timeout = None
        mock_trigger.run.side_effect = asyncio.CancelledError()

        with pytest.raises(asyncio.CancelledError):
            await trigger_runner.run_trigger(1, mock_trigger)

    @pytest.mark.asyncio
    @patch("airflow.jobs.triggerer_job_runner.TriggerRunner.set_individual_trigger_logging")
    async def test_run_trigger_timeout(self, session, caplog) -> None:
        trigger_runner = TriggerRunner()
        trigger_runner.triggers = {1: {"task": MagicMock(), "name": "mock_name", "events": 0}}
        mock_trigger = MagicMock()
        mock_trigger.task_instance.trigger_timeout = timezone.utcnow() - datetime.timedelta(hours=1)
        mock_trigger.run.side_effect = asyncio.CancelledError()

        with pytest.raises(asyncio.CancelledError):
            await trigger_runner.run_trigger(1, mock_trigger)
        assert "Trigger cancelled due to timeout" in caplog.text

    @patch("airflow.models.trigger.Trigger.bulk_fetch")
    @patch(
        "airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath",
        return_value=DateTimeTrigger,
    )
    def test_update_trigger_with_triggerer_argument_change(
        self, mock_bulk_fetch, mock_get_trigger_by_classpath, session, caplog
    ) -> None:
        trigger_runner = TriggerRunner()
        mock_trigger_orm = MagicMock()
        mock_trigger_orm.kwargs = {"moment": ..., "not_exists_arg": ...}
        mock_get_trigger_by_classpath.return_value = {1: mock_trigger_orm}

        trigger_runner.update_triggers({1})

        assert "Trigger failed" in caplog.text
        assert "got an unexpected keyword argument 'not_exists_arg'" in caplog.text


def test_trigger_create_race_condition_18392(session, tmp_path):
    """
    This verifies the resolution of race condition documented in github issue #18392.
    Triggers are queued for creation by TriggerJob.load_triggers.
    There was a race condition where multiple triggers would be created unnecessarily.
    What happens is the runner completes the trigger and purges from the "running" list.
    Then job.load_triggers is called and it looks like the trigger is not running but should,
    so it queues it again.

    The scenario is as follows:
        1. job.load_triggers (trigger now queued)
        2. runner.create_triggers (trigger now running)
        3. job.handle_events (trigger still appears running so state not updated in DB)
        4. runner.cleanup_finished_triggers (trigger completed at this point; trigger from "running" set)
        5. job.load_triggers (trigger not running, but also not purged from DB, so it is queued again)
        6. runner.create_triggers (trigger created again)

    This test verifies that under this scenario only one trigger is created.
    """
    path = tmp_path / "test_trigger_bad_respawn.txt"

    class TriggerRunner_(TriggerRunner):
        """We do some waiting for main thread looping"""

        async def wait_for_job_method_count(self, method, count):
            for _ in range(30):
                await asyncio.sleep(0.1)
                if getattr(self, f"{method}_count", 0) >= count:
                    break
            else:
                pytest.fail(f"did not observe count {count} in job method {method}")

        async def create_triggers(self):
            """
            On first run, wait for job.load_triggers to make sure they are queued
            """
            if getattr(self, "loop_count", 0) == 0:
                await self.wait_for_job_method_count("load_triggers", 1)
            await super().create_triggers()
            self.loop_count = getattr(self, "loop_count", 0) + 1

        async def cleanup_finished_triggers(self):
            """On loop 1, make sure that job.handle_events was already called"""
            if self.loop_count == 1:
                await self.wait_for_job_method_count("handle_events", 1)
            await super().cleanup_finished_triggers()

    class TriggererJob_(TriggererJobRunner):
        """We do some waiting for runner thread looping (and track calls in job thread)"""

        def wait_for_runner_loop(self, runner_loop_count):
            for _ in range(30):
                time.sleep(0.1)
                if getattr(self.trigger_runner, "call_count", 0) >= runner_loop_count:
                    break
            else:
                pytest.fail("did not observe 2 loops in the runner thread")

        def load_triggers(self):
            """On second run, make sure that runner has called create_triggers in its second loop"""
            super().load_triggers()
            self.trigger_runner.load_triggers_count = (
                getattr(self.trigger_runner, "load_triggers_count", 0) + 1
            )
            if self.trigger_runner.load_triggers_count == 2:
                self.wait_for_runner_loop(runner_loop_count=2)

        def handle_events(self):
            super().handle_events()
            self.trigger_runner.handle_events_count = (
                getattr(self.trigger_runner, "handle_events_count", 0) + 1
            )

    trigger = TimeDeltaTrigger_(delta=datetime.timedelta(microseconds=1), filename=path.as_posix())
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    session.add(trigger_orm)

    dag = DagModel(dag_id="test-dag")
    dag_run = DagRun(dag.dag_id, run_id="abc", run_type="none")
    ti = TaskInstance(PythonOperator(task_id="dummy-task", python_callable=print), run_id=dag_run.run_id)
    ti.dag_id = dag.dag_id
    ti.trigger_id = 1
    session.add(dag)
    session.add(dag_run)
    session.add(ti)

    session.commit()

    job = Job()
    job_runner = TriggererJob_(job)
    job_runner.trigger_runner = TriggerRunner_()
    thread = Thread(target=job_runner._execute)
    thread.start()
    try:
        for _ in range(40):
            time.sleep(0.1)
            # ready to evaluate after 2 loops
            if getattr(job_runner.trigger_runner, "loop_count", 0) >= 2:
                break
        else:
            pytest.fail("did not observe 2 loops in the runner thread")
    finally:
        job_runner.trigger_runner.stop = True
        job_runner.trigger_runner.join()
        thread.join()
    instances = path.read_text().splitlines()
    assert len(instances) == 1


def test_trigger_from_dead_triggerer(session, create_task_instance):
    """
    Checks that the triggerer will correctly claim a Trigger that is assigned to a
    triggerer that does not exist.
    """
    # Use a trigger that has an invalid triggerer_id
    trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    trigger_orm.triggerer_id = 999  # Non-existent triggerer
    session.add(trigger_orm)
    ti_orm = create_task_instance(
        task_id="ti_orm",
        execution_date=datetime.datetime.utcnow(),
        run_id="orm_run_id",
    )
    ti_orm.trigger_id = trigger_orm.id
    session.add(trigger_orm)
    session.commit()
    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()
    # Make sure it turned up in TriggerRunner's queue
    assert [x for x, y in job_runner.trigger_runner.to_create] == [1]


def test_trigger_from_expired_triggerer(session, create_task_instance):
    """
    Checks that the triggerer will correctly claim a Trigger that is assigned to a
    triggerer that has an expired heartbeat.
    """
    # Use a trigger assigned to the expired triggerer
    trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    trigger_orm.triggerer_id = 42
    session.add(trigger_orm)
    ti_orm = create_task_instance(
        task_id="ti_orm",
        execution_date=datetime.datetime.utcnow(),
        run_id="orm_run_id",
    )
    ti_orm.trigger_id = trigger_orm.id
    session.add(trigger_orm)
    # Use a TriggererJobRunner with an expired heartbeat
    triggerer_job_orm = Job(TriggererJobRunner.job_type)
    triggerer_job_orm.id = 42
    triggerer_job_orm.start_date = timezone.utcnow() - datetime.timedelta(hours=1)
    triggerer_job_orm.end_date = None
    triggerer_job_orm.latest_heartbeat = timezone.utcnow() - datetime.timedelta(hours=1)
    session.add(triggerer_job_orm)
    session.commit()
    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job(TriggererJobRunner.job_type)
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()
    # Make sure it turned up in TriggerRunner's queue
    assert [x for x, y in job_runner.trigger_runner.to_create] == [1]


def test_trigger_runner_exception_stops_triggerer(session):
    """
    Checks that if an exception occurs when creating triggers, that the triggerer
    process stops
    """

    class MockTriggerException(Exception):
        pass

    class TriggerRunner_(TriggerRunner):
        async def create_triggers(self):
            raise MockTriggerException("Trigger creation failed")

    # Use a trigger that will immediately succeed
    trigger = SuccessTrigger()
    create_trigger_in_db(session, trigger)

    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.trigger_runner = TriggerRunner_()
    thread = Thread(target=job_runner._execute)
    thread.start()

    # Wait 4 seconds for the triggerer to stop
    try:
        for _ in range(40):
            time.sleep(0.1)
            if not thread.is_alive():
                break
        else:
            pytest.fail("TriggererJobRunner did not stop after exception in TriggerRunner")

        if not job_runner.trigger_runner.stop:
            pytest.fail("TriggerRunner not marked as stopped after exception in TriggerRunner")

    finally:
        job_runner.trigger_runner.stop = True
        # with suppress(MockTriggerException):
        job_runner.trigger_runner.join()
        thread.join()


def test_trigger_firing(session):
    """
    Checks that when a trigger fires, it correctly makes it into the
    event queue.
    """
    # Use a trigger that will immediately succeed
    trigger = SuccessTrigger()
    create_trigger_in_db(session, trigger)
    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()
    # Now, start TriggerRunner up (and set it as a daemon thread during tests)
    job_runner.daemon = True
    job_runner.trigger_runner.start()
    try:
        # Wait for up to 3 seconds for it to fire and appear in the event queue
        for _ in range(30):
            if job_runner.trigger_runner.events:
                assert list(job_runner.trigger_runner.events) == [(1, TriggerEvent(True))]
                break
            time.sleep(0.1)
        else:
            pytest.fail("TriggerRunner never sent the trigger event out")
    finally:
        # We always have to stop the runner
        job_runner.trigger_runner.stop = True


def test_trigger_failing(session):
    """
    Checks that when a trigger fails, it correctly makes it into the
    failure queue.
    """
    # Use a trigger that will immediately fail
    trigger = FailureTrigger()
    create_trigger_in_db(session, trigger)
    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()
    # Now, start TriggerRunner up (and set it as a daemon thread during tests)
    job_runner.daemon = True
    job_runner.trigger_runner.start()
    try:
        # Wait for up to 3 seconds for it to fire and appear in the event queue
        for _ in range(30):
            if job_runner.trigger_runner.failed_triggers:
                assert len(job_runner.trigger_runner.failed_triggers) == 1
                trigger_id, exc = next(iter(job_runner.trigger_runner.failed_triggers))
                assert trigger_id == 1
                assert isinstance(exc, ValueError)
                assert exc.args[0] == "Deliberate trigger failure"
                break
            time.sleep(0.1)
        else:
            pytest.fail("TriggerRunner never marked the trigger as failed")
    finally:
        # We always have to stop the runner
        job_runner.trigger_runner.stop = True


def test_trigger_cleanup(session):
    """
    Checks that the triggerer will correctly clean up triggers that do not
    have any task instances depending on them.
    """
    # Use a trigger that will not fire for the lifetime of the test
    # (we want to avoid it firing and deleting itself)
    trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
    trigger_orm = Trigger.from_object(trigger)
    trigger_orm.id = 1
    session.add(trigger_orm)
    session.commit()
    # Trigger the cleanup code
    Trigger.clean_unused(session=session)
    session.commit()
    # Make sure it's gone
    assert session.query(Trigger).count() == 0


def test_invalid_trigger(session, dag_maker):
    """
    Checks that the triggerer will correctly fail task instances that depend on
    triggers that can't even be loaded.
    """
    # Create a totally invalid trigger
    trigger_orm = Trigger(classpath="fake.classpath", kwargs={})
    trigger_orm.id = 1
    session.add(trigger_orm)
    session.commit()

    # Create the test DAG and task
    with dag_maker(dag_id="test_invalid_trigger", session=session):
        EmptyOperator(task_id="dummy1")

    dr = dag_maker.create_dagrun()
    task_instance = dr.task_instances[0]
    # Make a task instance based on that and tie it to the trigger
    task_instance.state = TaskInstanceState.DEFERRED
    task_instance.trigger_id = 1
    session.commit()

    # Make a TriggererJobRunner and have it retrieve DB tasks
    job = Job()
    job_runner = TriggererJobRunner(job)
    job_runner.load_triggers()

    # Make sure it turned up in the failed queue
    assert len(job_runner.trigger_runner.failed_triggers) == 1

    # Run the failed trigger handler
    job_runner.handle_failed_triggers()

    # Make sure it marked the task instance as failed (which is actually the
    # scheduled state with a payload to make it fail)
    task_instance.refresh_from_db()
    assert task_instance.state == TaskInstanceState.SCHEDULED
    assert task_instance.next_method == "__fail__"
    assert task_instance.next_kwargs["error"] == "Trigger failure"
    assert task_instance.next_kwargs["traceback"][-1] == "ModuleNotFoundError: No module named 'fake'\n"


@pytest.mark.parametrize("should_wrap", (True, False))
@patch("airflow.jobs.triggerer_job_runner.configure_trigger_log_handler")
def test_handler_config_respects_donot_wrap(mock_configure, should_wrap):
    from airflow.jobs import triggerer_job_runner

    triggerer_job_runner.DISABLE_WRAPPER = not should_wrap
    job = Job()
    TriggererJobRunner(job=job)
    if should_wrap:
        mock_configure.assert_called()
    else:
        mock_configure.assert_not_called()


@patch("airflow.jobs.triggerer_job_runner.setup_queue_listener")
def test_triggerer_job_always_creates_listener(mock_setup):
    mock_setup.assert_not_called()
    job = Job()
    TriggererJobRunner(job=job)
    mock_setup.assert_called()


def test_queue_listener():
    """
    When listener func called, root handlers should be moved to queue listener
    and replaced with queuehandler.
    """
    reset_logging()
    importlib.reload(airflow_local_settings)
    configure_logging()

    def non_pytest_handlers(val):
        return [h for h in val if "pytest" not in h.__module__]

    import logging

    log = logging.getLogger()
    handlers = non_pytest_handlers(log.handlers)
    assert len(handlers) == 1
    handler = handlers[0]
    assert handler.__class__ == RedirectStdHandler
    listener = setup_queue_listener()
    assert handler not in non_pytest_handlers(log.handlers)
    qh = log.handlers[-1]
    assert qh.__class__ == LocalQueueHandler
    assert qh.queue == listener.queue
    listener.stop()
