# 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 logging
import os
from datetime import datetime, timedelta
from pathlib import Path
from unittest import mock
from unittest.mock import MagicMock, patch
from uuid import uuid4

import pytest
import time_machine
from sqlalchemy import delete, select

from airflow.executors.workloads import BundleInfo, ExecuteTask
from airflow.providers.common.compat.sdk import Stats, TaskInstanceKey, conf, timezone
from airflow.providers.edge3.executors.edge_executor import EdgeExecutor
from airflow.providers.edge3.models.edge_job import EdgeJobModel
from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState
from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG
from airflow.utils.session import create_session
from airflow.utils.state import TaskInstanceState

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS, AIRFLOW_V_3_3_PLUS

if AIRFLOW_V_3_3_PLUS:
    from airflow.executors.workloads import CallbackFetchMethod, ExecuteCallback, TaskInstanceDTO
    from airflow.executors.workloads.callback import CallbackDTO

pytestmark = pytest.mark.db_test


class TestEdgeExecutor:
    @pytest.fixture(autouse=True)
    def setup_test_cases(self):
        with create_session() as session:
            session.execute(delete(EdgeJobModel))

    def get_test_executor(self, pool_slots=1):
        key = TaskInstanceKey(
            dag_id="test_dag", run_id="test_run", task_id="test_task", map_index=-1, try_number=1
        )
        ti = MagicMock()
        ti.pool_slots = pool_slots
        ti.dag_run.dag_id = key.dag_id
        ti.dag_run.run_id = key.run_id
        ti.dag_run.start_date = datetime(2021, 1, 1)
        executor = EdgeExecutor()
        executor.queued_tasks = {key: [None, None, None, ti]}

        return (executor, key)

    @patch(f"{Stats.__module__}.Stats.incr")
    def test_sync_orphaned_tasks(self, mock_stats_incr):
        executor = EdgeExecutor()

        delta_to_purge = timedelta(minutes=conf.getint("edge", "job_fail_purge") + 1)
        delta_to_orphaned_config_name = "task_instance_heartbeat_timeout"

        delta_to_orphaned = timedelta(seconds=conf.getint("scheduler", delta_to_orphaned_config_name) + 1)

        with create_session() as session:
            for task_id, state, last_update in [
                (
                    "started_running_orphaned",
                    TaskInstanceState.RUNNING,
                    timezone.utcnow() - delta_to_orphaned,
                ),
                ("started_removed", TaskInstanceState.REMOVED, timezone.utcnow() - delta_to_purge),
            ]:
                session.add(
                    EdgeJobModel(
                        dag_id="test_dag",
                        task_id=task_id,
                        run_id="test_run",
                        map_index=-1,
                        try_number=1,
                        state=state,
                        queue="default",
                        command="mock",
                        concurrency_slots=1,
                        last_update=last_update,
                    )
                )
                session.commit()

        executor.sync()

        mock_stats_incr.assert_called_with(
            "edge_worker.ti.finish",
            tags={
                "dag_id": "test_dag",
                "queue": "default",
                "state": "failed",
                "task_id": "started_running_orphaned",
            },
        )
        assert mock_stats_incr.call_count == 1

        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(jobs) == 1
            assert jobs[0].task_id == "started_running_orphaned"
            assert jobs[0].state == TaskInstanceState.REMOVED

    @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.running_state")
    @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.success")
    @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.fail")
    def test_sync(self, mock_fail, mock_success, mock_running_state):
        executor = EdgeExecutor()

        def remove_from_running(key: TaskInstanceKey):
            executor.running.remove(key)

        mock_success.side_effect = remove_from_running
        mock_fail.side_effect = remove_from_running

        delta_to_purge = timedelta(minutes=conf.getint("edge", "job_fail_purge") + 1)

        # Prepare some data
        with create_session() as session:
            for task_id, state, last_update in [
                ("started_running", TaskInstanceState.RUNNING, timezone.utcnow()),
                ("started_success", TaskInstanceState.SUCCESS, timezone.utcnow() - delta_to_purge),
                ("started_failed", TaskInstanceState.FAILED, timezone.utcnow() - delta_to_purge),
            ]:
                session.add(
                    EdgeJobModel(
                        dag_id="test_dag",
                        task_id=task_id,
                        run_id="test_run",
                        map_index=-1,
                        try_number=1,
                        state=state,
                        queue="default",
                        concurrency_slots=1,
                        command="mock",
                        last_update=last_update,
                    )
                )
                key = TaskInstanceKey(
                    dag_id="test_dag", run_id="test_run", task_id=task_id, map_index=-1, try_number=1
                )
                executor.running.add(key)
                session.commit()
        assert len(executor.running) == 3

        executor.sync()

        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(session.scalars(select(EdgeJobModel)).all()) == 1
            assert jobs[0].task_id == "started_running"
            assert jobs[0].state == TaskInstanceState.RUNNING

        assert len(executor.running) == 1
        mock_running_state.assert_called_once()
        mock_success.assert_called_once()
        mock_fail.assert_called_once()

        # Any test another round with one new run
        mock_running_state.reset_mock()
        mock_success.reset_mock()
        mock_fail.reset_mock()

        with create_session() as session:
            task_id = "started_running2"
            state = TaskInstanceState.RUNNING
            session.add(
                EdgeJobModel(
                    dag_id="test_dag",
                    task_id=task_id,
                    run_id="test_run",
                    map_index=-1,
                    try_number=1,
                    state=state,
                    concurrency_slots=1,
                    queue="default",
                    command="mock",
                    last_update=timezone.utcnow(),
                )
            )
            key = TaskInstanceKey(
                dag_id="test_dag", run_id="test_run", task_id=task_id, map_index=-1, try_number=1
            )
            executor.running.add(key)
            session.commit()
        assert len(executor.running) == 2

        executor.sync()

        assert len(executor.running) == 2
        mock_running_state.assert_called_once()  # because we reported already first run, new run calling
        mock_success.assert_not_called()  # because we reported already, not running anymore
        mock_fail.assert_not_called()  # because we reported already, not running anymore
        mock_running_state.reset_mock()

        executor.sync()

        assert len(executor.running) == 2
        # Now none is called as we called before already
        mock_running_state.assert_not_called()
        mock_success.assert_not_called()
        mock_fail.assert_not_called()

    def test_sync_active_worker(self):
        executor = EdgeExecutor()

        # Prepare some data
        with create_session() as session:
            # Clear existing workers to avoid unique constraint violation
            session.execute(delete(EdgeWorkerModel))
            session.commit()

            # Add workers with different states
            for worker_name, state, last_heartbeat in [
                (
                    "inactive_timed_out_worker",
                    EdgeWorkerState.IDLE,
                    datetime(2023, 1, 1, 0, 59, 0, tzinfo=timezone.utc),
                ),
                ("active_worker", EdgeWorkerState.IDLE, datetime(2023, 1, 1, 0, 59, 10, tzinfo=timezone.utc)),
                (
                    "offline_worker",
                    EdgeWorkerState.OFFLINE,
                    datetime(2023, 1, 1, 0, 59, 10, tzinfo=timezone.utc),
                ),
                (
                    "offline_maintenance_worker",
                    EdgeWorkerState.OFFLINE_MAINTENANCE,
                    datetime(2023, 1, 1, 0, 59, 10, tzinfo=timezone.utc),
                ),
            ]:
                ewm = EdgeWorkerModel(
                    worker_name=worker_name,
                    state=state,
                    last_update=last_heartbeat,
                    queues="",
                    first_online=timezone.utcnow(),
                )
                ewm.sysinfo = {"status": logging.INFO, "status_text": "I am good, sun is shining 🌞"}
                session.add(ewm)
                session.commit()

        with time_machine.travel(datetime(2023, 1, 1, 1, 0, 0, tzinfo=timezone.utc), tick=False):
            with conf_vars({("edge", "heartbeat_interval"): "10"}):
                executor.sync()

        with create_session() as session:
            for worker in session.scalars(select(EdgeWorkerModel)).all():
                print(worker.worker_name)
                if "maintenance_" in worker.worker_name:
                    assert worker.state == EdgeWorkerState.OFFLINE_MAINTENANCE
                elif "offline_" in worker.worker_name:
                    assert worker.state == EdgeWorkerState.OFFLINE
                elif "inactive_" in worker.worker_name:
                    assert worker.state == EdgeWorkerState.UNKNOWN
                    assert worker.sysinfo
                    assert worker.sysinfo["status"] == logging.NOTSET
                    assert "status_text" not in worker.sysinfo
                else:
                    assert worker.state == EdgeWorkerState.IDLE
                    assert worker.sysinfo
                    assert worker.sysinfo["status"] == logging.INFO
                    assert "status_text" in worker.sysinfo

    def test_revoke_task(self):
        """Test that revoke_task removes task from executor and database."""
        executor = EdgeExecutor()
        key = TaskInstanceKey(
            dag_id="test_dag", run_id="test_run", task_id="test_task", map_index=-1, try_number=1
        )

        # Create a mock task instance
        ti = MagicMock()
        ti.key = key
        ti.dag_id = "test_dag"
        ti.task_id = "test_task"
        ti.run_id = "test_run"
        ti.map_index = -1
        ti.try_number = 1

        # Add task to executor's internal state
        executor.running.add(key)
        executor.queued_tasks[key] = [None, None, None, ti]
        executor.last_reported_state[key] = TaskInstanceState.QUEUED

        # Add corresponding job to database
        with create_session() as session:
            session.add(
                EdgeJobModel(
                    dag_id="test_dag",
                    task_id="test_task",
                    run_id="test_run",
                    map_index=-1,
                    try_number=1,
                    state=TaskInstanceState.QUEUED,
                    queue="default",
                    command="mock",
                    concurrency_slots=1,
                )
            )
            session.commit()

        # Verify job exists before revoke
        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(jobs) == 1

        # Revoke the task
        executor.revoke_task(ti=ti)

        # Verify task is removed from executor's internal state
        assert key not in executor.running
        assert key not in executor.queued_tasks
        assert key not in executor.last_reported_state

        # Verify job is removed from database
        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(jobs) == 0

    def test_revoke_task_nonexistent(self):
        """Test that revoke_task handles non-existent tasks gracefully."""
        executor = EdgeExecutor()
        key = TaskInstanceKey(
            dag_id="nonexistent_dag",
            run_id="nonexistent_run",
            task_id="nonexistent_task",
            map_index=-1,
            try_number=1,
        )

        # Create a mock task instance
        ti = MagicMock()
        ti.key = key
        ti.dag_id = "nonexistent_dag"
        ti.task_id = "nonexistent_task"
        ti.run_id = "nonexistent_run"
        ti.map_index = -1
        ti.try_number = 1

        # Revoke a task that doesn't exist (should not raise error)
        executor.revoke_task(ti=ti)

        # Verify nothing breaks
        assert key not in executor.running
        assert key not in executor.queued_tasks


class TestEdgeExecutorMultiTeam:
    """Tests for multi-team (AIP-67) support in EdgeExecutor."""

    @pytest.fixture(autouse=True)
    def setup_test_cases(self):
        with create_session() as session:
            session.execute(delete(EdgeJobModel))
            session.execute(delete(EdgeWorkerModel))
            session.commit()

    def test_global_executor_without_team_name(self):
        """Test that global executor (no team) works correctly."""
        executor = EdgeExecutor()

        assert hasattr(executor, "conf")
        assert executor.team_name is None
        if AIRFLOW_V_3_2_PLUS:
            assert executor.conf.team_name is None

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_executor_with_team_name(self):
        """Test that executor with team_name has correct conf setup."""
        team_name = "test_team"
        executor = EdgeExecutor(team_name=team_name)

        assert hasattr(executor, "conf")
        assert executor.team_name == team_name
        assert executor.conf.team_name == team_name

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_multiple_team_executors_isolation(self):
        """Test that multiple team executors can coexist with isolated resources."""
        team_a_executor = EdgeExecutor(parallelism=2, team_name="team_a")
        team_b_executor = EdgeExecutor(parallelism=3, team_name="team_b")

        assert team_a_executor.running is not team_b_executor.running
        assert team_a_executor.queued_tasks is not team_b_executor.queued_tasks
        assert team_a_executor.last_reported_state is not team_b_executor.last_reported_state

        if AIRFLOW_V_3_2_PLUS:
            assert team_a_executor.conf.team_name == "team_a"
            assert team_b_executor.conf.team_name == "team_b"

        assert team_a_executor.parallelism == 2
        assert team_b_executor.parallelism == 3

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_team_config_used_in_check_worker_liveness(self):
        """Test that _check_worker_liveness reads config from self.conf, not global conf."""
        team_name = "test_team"
        executor = EdgeExecutor(team_name=team_name)

        team_env_key_prefix = f"AIRFLOW__{team_name.upper()}___EDGE__"
        test_key_values = [
            "heartbeat_interval",
            "task_instance_heartbeat_timeout",
            "job_success_purge",
            "job_fail_purge",
        ]
        for test_key_value in test_key_values:
            with mock.patch.dict(os.environ, {f"{team_env_key_prefix}{test_key_value.upper()}": "100"}):
                value = executor.conf.getint("edge", test_key_value)
                assert value == 100

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_purge_jobs_filters_by_team_name(self):
        """Test that _purge_jobs only purges jobs belonging to its team."""
        executor_a = EdgeExecutor(team_name="team_a")

        delta_to_purge = timedelta(minutes=conf.getint("edge", "job_fail_purge") + 1)

        with create_session() as session:
            for team in ["team_a", "team_b"]:
                session.add(
                    EdgeJobModel(
                        dag_id="test_dag",
                        task_id=f"task_{team}",
                        run_id="test_run",
                        map_index=-1,
                        try_number=1,
                        state=TaskInstanceState.FAILED,
                        queue="default",
                        command="mock",
                        concurrency_slots=1,
                        last_update=timezone.utcnow() - delta_to_purge,
                        team_name=team,
                    )
                )
            session.commit()

        with create_session() as session:
            executor_a._purge_jobs(session)
            session.commit()

        with create_session() as session:
            remaining_jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(remaining_jobs) == 1
            assert remaining_jobs[0].team_name == "team_b"

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_update_orphaned_jobs_filters_by_team_name(self):
        """Test that _update_orphaned_jobs only checks jobs belonging to its team."""
        executor_a = EdgeExecutor(team_name="team_a")

        heartbeat_timeout = conf.getint("scheduler", "task_instance_heartbeat_timeout")
        delta_to_orphaned = timedelta(seconds=heartbeat_timeout + 1)

        with create_session() as session:
            for team in ["team_a", "team_b"]:
                session.add(
                    EdgeJobModel(
                        dag_id="test_dag",
                        task_id=f"task_{team}",
                        run_id="test_run",
                        map_index=-1,
                        try_number=1,
                        state=TaskInstanceState.RUNNING,
                        queue="default",
                        command="mock",
                        concurrency_slots=1,
                        last_update=timezone.utcnow() - delta_to_orphaned,
                        team_name=team,
                    )
                )
            session.commit()

        with create_session() as session:
            executor_a._update_orphaned_jobs(session)
            session.commit()

        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            jobs_by_team = {job.team_name: job for job in jobs}
            assert jobs_by_team["team_a"].state != TaskInstanceState.RUNNING
            assert jobs_by_team["team_b"].state == TaskInstanceState.RUNNING

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_check_worker_liveness_filters_by_team_name(self):
        """Test that _check_worker_liveness only resets workers belonging to its team."""
        executor_a = EdgeExecutor(team_name="team_a")

        with create_session() as session:
            for worker_name, team in [
                ("worker_team_a", "team_a"),
                ("worker_team_b", "team_b"),
            ]:
                session.add(
                    EdgeWorkerModel(
                        worker_name=worker_name,
                        state=EdgeWorkerState.IDLE,
                        last_update=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
                        queues=["default"],
                        first_online=timezone.utcnow(),
                        team_name=team,
                    )
                )
            session.commit()

        with time_machine.travel(datetime(2023, 1, 1, 1, 0, 0, tzinfo=timezone.utc), tick=False):
            with conf_vars({("edge", "heartbeat_interval"): "10"}):
                with create_session() as session:
                    executor_a._check_worker_liveness(session)
                    session.commit()

        with create_session() as session:
            workers = {w.worker_name: w for w in session.scalars(select(EdgeWorkerModel)).all()}
            assert workers["worker_team_a"].state == EdgeWorkerState.UNKNOWN
            assert workers["worker_team_b"].state == EdgeWorkerState.IDLE

    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="The tests should be skipped for Airflow < 3.2")
    def test_no_team_executor_processes_all_jobs(self):
        """Test that an executor without team_name processes only job has no team_name."""
        executor = EdgeExecutor()

        delta_to_purge = timedelta(minutes=conf.getint("edge", "job_fail_purge") + 1)

        with create_session() as session:
            for team in ["team_a", "team_b", None]:
                session.add(
                    EdgeJobModel(
                        dag_id="test_dag",
                        task_id=f"task_{team}",
                        run_id="test_run",
                        map_index=-1,
                        try_number=1,
                        state=TaskInstanceState.FAILED,
                        queue="default",
                        command="mock",
                        concurrency_slots=1,
                        last_update=timezone.utcnow() - delta_to_purge,
                        team_name=team,
                    )
                )
            session.commit()

        with create_session() as session:
            executor._purge_jobs(session)
            session.commit()

        with create_session() as session:
            remaining_jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(remaining_jobs) == 2


@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="ExecuteTypeBody union requires Airflow 3.3+")
class TestQueueWorkload:
    @pytest.fixture(autouse=True)
    def setup(self):
        with create_session() as session:
            session.execute(delete(EdgeJobModel))
            session.commit()

    def _make_execute_task(self) -> ExecuteTask:
        ti = TaskInstanceDTO(
            id=uuid4(),
            dag_version_id=uuid4(),
            task_id="test_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            map_index=-1,
            pool_slots=1,
            queue="default",
            priority_weight=1,
        )
        return ExecuteTask(
            ti=ti,
            dag_rel_path=Path("test_dag.py"),
            token="test_token",
            bundle_info=BundleInfo(name="test_bundle", version="1.0"),
            log_path="test.log",
        )

    def test_queue_workload_execute_task(self):
        executor = EdgeExecutor()
        workload = self._make_execute_task()

        with create_session() as session:
            executor.queue_workload(workload, session=session)

        with create_session() as session:
            job = session.scalar(select(EdgeJobModel))
            assert job is not None
            assert job.dag_id == "test_dag"
            assert job.task_id == "test_task"
            assert job.run_id == "test_run"
            assert job.state == TaskInstanceState.QUEUED
            assert '"type":"ExecuteTask"' in job.command or '"type": "ExecuteTask"' in job.command

    def test_queue_workload_execute_task_existing_job(self):
        executor = EdgeExecutor()
        workload = self._make_execute_task()

        with create_session() as session:
            executor.queue_workload(workload, session=session)
        with create_session() as session:
            executor.queue_workload(workload, session=session)

        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(jobs) == 1
            assert jobs[0].state == TaskInstanceState.QUEUED

    def test_queue_workload_execute_callback(self):
        executor = EdgeExecutor()
        id = str(uuid4())
        callback_data = CallbackDTO(
            id=id,
            fetch_method=CallbackFetchMethod.IMPORT_PATH,
            data={
                "path": "builtins.dict",
                "kwargs": {"a": 1, "b": 2, "c": 3},
            },
        )
        workload = ExecuteCallback(
            callback=callback_data,
            dag_rel_path=Path("test.py"),
            bundle_info=BundleInfo(name="test_bundle", version="1.0"),
            token="test_token",
            log_path="test.log",
        )

        with create_session() as session:
            executor.queue_workload(workload, session=session)

        with create_session() as session:
            job = session.scalar(select(EdgeJobModel))
            assert job is not None
            assert job.dag_id == EXECUTE_CALLBACK_TAG
            assert job.task_id == id
            assert job.run_id == f"{EXECUTE_CALLBACK_TAG}-{id}"
            assert job.state == TaskInstanceState.QUEUED
            assert '"type":"ExecuteCallback"' in job.command or '"type": "ExecuteCallback"' in job.command

    def test_queue_workload_execute_callback_existing_job(self):
        executor = EdgeExecutor()
        callback_data = CallbackDTO(
            id=str(uuid4()),
            fetch_method=CallbackFetchMethod.IMPORT_PATH,
            data={
                "path": "builtins.dict",
                "kwargs": {"a": 1, "b": 2, "c": 3},
            },
        )
        workload = ExecuteCallback(
            callback=callback_data,
            dag_rel_path=Path("test.py"),
            bundle_info=BundleInfo(name="test_bundle", version="1.0"),
            token="test_token",
            log_path="test.log",
        )

        with create_session() as session:
            executor.queue_workload(workload, session=session)
        with create_session() as session:
            executor.queue_workload(workload, session=session)

        with create_session() as session:
            jobs = session.scalars(select(EdgeJobModel)).all()
            assert len(jobs) == 1
            assert jobs[0].state == TaskInstanceState.QUEUED

    def test_queue_workload_unknown_type_raises(self):
        executor = EdgeExecutor()
        with create_session() as session:
            with pytest.raises(TypeError, match="Don't know how to queue workload"):
                executor.queue_workload(MagicMock(spec=[]), session=session)
