# 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 os
import unittest.mock
from datetime import datetime

import pytest

from airflow import DAG
from airflow.models import DagBag
from airflow.models.expandinput import EXPAND_INPUT_EMPTY
from airflow.models.serialized_dag import SerializedDagModel
from airflow.operators.empty import EmptyOperator
from airflow.security import permissions
from tests.test_utils.api_connexion_utils import assert_401, create_user, delete_user
from tests.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags


@pytest.fixture(scope="module")
def configured_app(minimal_app_for_api):
    app = minimal_app_for_api
    create_user(
        app,  # type: ignore
        username="test",
        role_name="Test",
        permissions=[
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG_RUN),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_TASK_INSTANCE),
        ],
    )
    create_user(app, username="test_no_permissions", role_name="TestNoPermissions")  # type: ignore

    yield app

    delete_user(app, username="test")  # type: ignore
    delete_user(app, username="test_no_permissions")  # type: ignore


class TestTaskEndpoint:
    dag_id = "test_dag"
    mapped_dag_id = "test_mapped_task"
    task_id = "op1"
    task_id2 = "op2"
    task_id3 = "op3"
    mapped_task_id = "mapped_task"
    task1_start_date = datetime(2020, 6, 15)
    task2_start_date = datetime(2020, 6, 16)

    @pytest.fixture(scope="class")
    def setup_dag(self, configured_app):
        with DAG(self.dag_id, start_date=self.task1_start_date, doc_md="details") as dag:
            task1 = EmptyOperator(task_id=self.task_id, params={"foo": "bar"})
            task2 = EmptyOperator(task_id=self.task_id2, start_date=self.task2_start_date)

        with DAG(self.mapped_dag_id, start_date=self.task1_start_date) as mapped_dag:
            EmptyOperator(task_id=self.task_id3)
            # Use the private _expand() method to avoid the empty kwargs check.
            # We don't care about how the operator runs here, only its presence.
            EmptyOperator.partial(task_id=self.mapped_task_id)._expand(EXPAND_INPUT_EMPTY, strict=False)

        task1 >> task2
        dag_bag = DagBag(os.devnull, include_examples=False)
        dag_bag.dags = {dag.dag_id: dag, mapped_dag.dag_id: mapped_dag}
        configured_app.dag_bag = dag_bag  # type:ignore

    @staticmethod
    def clean_db():
        clear_db_runs()
        clear_db_dags()
        clear_db_serialized_dags()

    @pytest.fixture(autouse=True)
    def setup_attrs(self, configured_app, setup_dag) -> None:
        self.clean_db()
        self.app = configured_app
        self.client = self.app.test_client()  # type:ignore

    def teardown_method(self) -> None:
        self.clean_db()


class TestGetTask(TestTaskEndpoint):
    def test_should_respond_200(self):
        expected = {
            "class_ref": {
                "class_name": "EmptyOperator",
                "module_path": "airflow.operators.empty",
            },
            "depends_on_past": False,
            "downstream_task_ids": [self.task_id2],
            "end_date": None,
            "execution_timeout": None,
            "extra_links": [],
            "operator_name": "EmptyOperator",
            "owner": "airflow",
            "params": {
                "foo": {
                    "__class": "airflow.models.param.Param",
                    "value": "bar",
                    "description": None,
                    "schema": {},
                }
            },
            "pool": "default_pool",
            "pool_slots": 1.0,
            "priority_weight": 1.0,
            "queue": "default",
            "retries": 0.0,
            "retry_delay": {"__type": "TimeDelta", "days": 0, "seconds": 300, "microseconds": 0},
            "retry_exponential_backoff": False,
            "start_date": "2020-06-15T00:00:00+00:00",
            "task_id": "op1",
            "template_fields": [],
            "trigger_rule": "all_success",
            "ui_color": "#e8f7e4",
            "ui_fgcolor": "#000",
            "wait_for_downstream": False,
            "weight_rule": "downstream",
            "is_mapped": False,
        }
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks/{self.task_id}", environ_overrides={"REMOTE_USER": "test"}
        )
        assert response.status_code == 200
        assert response.json == expected

    def test_mapped_task(self):
        expected = {
            "class_ref": {"class_name": "EmptyOperator", "module_path": "airflow.operators.empty"},
            "depends_on_past": False,
            "downstream_task_ids": [],
            "end_date": None,
            "execution_timeout": None,
            "extra_links": [],
            "is_mapped": True,
            "operator_name": "EmptyOperator",
            "owner": "airflow",
            "params": {},
            "pool": "default_pool",
            "pool_slots": 1.0,
            "priority_weight": 1.0,
            "queue": "default",
            "retries": 0.0,
            "retry_delay": {"__type": "TimeDelta", "days": 0, "microseconds": 0, "seconds": 300},
            "retry_exponential_backoff": False,
            "start_date": "2020-06-15T00:00:00+00:00",
            "task_id": "mapped_task",
            "template_fields": [],
            "trigger_rule": "all_success",
            "ui_color": "#e8f7e4",
            "ui_fgcolor": "#000",
            "wait_for_downstream": False,
            "weight_rule": "downstream",
        }
        response = self.client.get(
            f"/api/v1/dags/{self.mapped_dag_id}/tasks/{self.mapped_task_id}",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200
        assert response.json == expected

    def test_should_respond_200_serialized(self):

        # Get the dag out of the dagbag before we patch it to an empty one
        SerializedDagModel.write_dag(self.app.dag_bag.get_dag(self.dag_id))

        dag_bag = DagBag(os.devnull, include_examples=False, read_dags_from_db=True)
        patcher = unittest.mock.patch.object(self.app, "dag_bag", dag_bag)
        patcher.start()

        expected = {
            "class_ref": {
                "class_name": "EmptyOperator",
                "module_path": "airflow.operators.empty",
            },
            "depends_on_past": False,
            "downstream_task_ids": [self.task_id2],
            "end_date": None,
            "execution_timeout": None,
            "extra_links": [],
            "operator_name": "EmptyOperator",
            "owner": "airflow",
            "params": {
                "foo": {
                    "__class": "airflow.models.param.Param",
                    "value": "bar",
                    "description": None,
                    "schema": {},
                }
            },
            "pool": "default_pool",
            "pool_slots": 1.0,
            "priority_weight": 1.0,
            "queue": "default",
            "retries": 0.0,
            "retry_delay": {"__type": "TimeDelta", "days": 0, "seconds": 300, "microseconds": 0},
            "retry_exponential_backoff": False,
            "start_date": "2020-06-15T00:00:00+00:00",
            "task_id": "op1",
            "template_fields": [],
            "trigger_rule": "all_success",
            "ui_color": "#e8f7e4",
            "ui_fgcolor": "#000",
            "wait_for_downstream": False,
            "weight_rule": "downstream",
            "is_mapped": False,
        }
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks/{self.task_id}", environ_overrides={"REMOTE_USER": "test"}
        )
        assert response.status_code == 200
        assert response.json == expected
        patcher.stop()

    def test_should_respond_404(self):
        task_id = "xxxx_not_existing"
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks/{task_id}", environ_overrides={"REMOTE_USER": "test"}
        )
        assert response.status_code == 404

    def test_should_raises_401_unauthenticated(self):
        response = self.client.get(f"/api/v1/dags/{self.dag_id}/tasks/{self.task_id}")

        assert_401(response)

    def test_should_raise_403_forbidden(self):
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks", environ_overrides={"REMOTE_USER": "test_no_permissions"}
        )
        assert response.status_code == 403


class TestGetTasks(TestTaskEndpoint):
    def test_should_respond_200(self):
        expected = {
            "tasks": [
                {
                    "class_ref": {
                        "class_name": "EmptyOperator",
                        "module_path": "airflow.operators.empty",
                    },
                    "depends_on_past": False,
                    "downstream_task_ids": [self.task_id2],
                    "end_date": None,
                    "execution_timeout": None,
                    "extra_links": [],
                    "operator_name": "EmptyOperator",
                    "owner": "airflow",
                    "params": {
                        "foo": {
                            "__class": "airflow.models.param.Param",
                            "value": "bar",
                            "description": None,
                            "schema": {},
                        }
                    },
                    "pool": "default_pool",
                    "pool_slots": 1.0,
                    "priority_weight": 1.0,
                    "queue": "default",
                    "retries": 0.0,
                    "retry_delay": {"__type": "TimeDelta", "days": 0, "seconds": 300, "microseconds": 0},
                    "retry_exponential_backoff": False,
                    "start_date": "2020-06-15T00:00:00+00:00",
                    "task_id": "op1",
                    "template_fields": [],
                    "trigger_rule": "all_success",
                    "ui_color": "#e8f7e4",
                    "ui_fgcolor": "#000",
                    "wait_for_downstream": False,
                    "weight_rule": "downstream",
                    "is_mapped": False,
                },
                {
                    "class_ref": {
                        "class_name": "EmptyOperator",
                        "module_path": "airflow.operators.empty",
                    },
                    "depends_on_past": False,
                    "downstream_task_ids": [],
                    "end_date": None,
                    "execution_timeout": None,
                    "extra_links": [],
                    "operator_name": "EmptyOperator",
                    "owner": "airflow",
                    "params": {},
                    "pool": "default_pool",
                    "pool_slots": 1.0,
                    "priority_weight": 1.0,
                    "queue": "default",
                    "retries": 0.0,
                    "retry_delay": {"__type": "TimeDelta", "days": 0, "seconds": 300, "microseconds": 0},
                    "retry_exponential_backoff": False,
                    "start_date": "2020-06-16T00:00:00+00:00",
                    "task_id": self.task_id2,
                    "template_fields": [],
                    "trigger_rule": "all_success",
                    "ui_color": "#e8f7e4",
                    "ui_fgcolor": "#000",
                    "wait_for_downstream": False,
                    "weight_rule": "downstream",
                    "is_mapped": False,
                },
            ],
            "total_entries": 2,
        }
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks", environ_overrides={"REMOTE_USER": "test"}
        )
        assert response.status_code == 200
        assert response.json == expected

    def test_get_tasks_mapped(self):
        expected = {
            "tasks": [
                {
                    "class_ref": {"class_name": "EmptyOperator", "module_path": "airflow.operators.empty"},
                    "depends_on_past": False,
                    "downstream_task_ids": [],
                    "end_date": None,
                    "execution_timeout": None,
                    "extra_links": [],
                    "is_mapped": True,
                    "operator_name": "EmptyOperator",
                    "owner": "airflow",
                    "params": {},
                    "pool": "default_pool",
                    "pool_slots": 1.0,
                    "priority_weight": 1.0,
                    "queue": "default",
                    "retries": 0.0,
                    "retry_delay": {"__type": "TimeDelta", "days": 0, "microseconds": 0, "seconds": 300},
                    "retry_exponential_backoff": False,
                    "start_date": "2020-06-15T00:00:00+00:00",
                    "task_id": "mapped_task",
                    "template_fields": [],
                    "trigger_rule": "all_success",
                    "ui_color": "#e8f7e4",
                    "ui_fgcolor": "#000",
                    "wait_for_downstream": False,
                    "weight_rule": "downstream",
                },
                {
                    "class_ref": {
                        "class_name": "EmptyOperator",
                        "module_path": "airflow.operators.empty",
                    },
                    "depends_on_past": False,
                    "downstream_task_ids": [],
                    "end_date": None,
                    "execution_timeout": None,
                    "extra_links": [],
                    "operator_name": "EmptyOperator",
                    "owner": "airflow",
                    "params": {},
                    "pool": "default_pool",
                    "pool_slots": 1.0,
                    "priority_weight": 1.0,
                    "queue": "default",
                    "retries": 0.0,
                    "retry_delay": {"__type": "TimeDelta", "days": 0, "seconds": 300, "microseconds": 0},
                    "retry_exponential_backoff": False,
                    "start_date": "2020-06-15T00:00:00+00:00",
                    "task_id": self.task_id3,
                    "template_fields": [],
                    "trigger_rule": "all_success",
                    "ui_color": "#e8f7e4",
                    "ui_fgcolor": "#000",
                    "wait_for_downstream": False,
                    "weight_rule": "downstream",
                    "is_mapped": False,
                },
            ],
            "total_entries": 2,
        }
        response = self.client.get(
            f"/api/v1/dags/{self.mapped_dag_id}/tasks", environ_overrides={"REMOTE_USER": "test"}
        )
        assert response.status_code == 200
        assert response.json == expected

    def test_should_respond_200_ascending_order_by_start_date(self):
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks?order_by=start_date",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200
        assert self.task1_start_date < self.task2_start_date
        assert response.json["tasks"][0]["task_id"] == self.task_id
        assert response.json["tasks"][1]["task_id"] == self.task_id2

    def test_should_respond_200_descending_order_by_start_date(self):
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks?order_by=-start_date",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200
        # - means is descending
        assert self.task1_start_date < self.task2_start_date
        assert response.json["tasks"][0]["task_id"] == self.task_id2
        assert response.json["tasks"][1]["task_id"] == self.task_id

    def test_should_raise_400_for_invalid_order_by_name(self):
        response = self.client.get(
            f"/api/v1/dags/{self.dag_id}/tasks?order_by=invalid_task_colume_name",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 400
        assert response.json["detail"] == "'EmptyOperator' object has no attribute 'invalid_task_colume_name'"

    def test_should_respond_404(self):
        dag_id = "xxxx_not_existing"
        response = self.client.get(f"/api/v1/dags/{dag_id}/tasks", environ_overrides={"REMOTE_USER": "test"})
        assert response.status_code == 404

    def test_should_raises_401_unauthenticated(self):
        response = self.client.get(f"/api/v1/dags/{self.dag_id}/tasks")

        assert_401(response)