"""
# Data Refresh TaskGroup Factory
This file generates the data refresh TaskGroup using a factory function.
This TaskGroup initiates a data refresh for a given media type and awaits the
success or failure of the refresh. Importantly, it is also configured to
ensure that no two remote data refreshes can run concurrently, as required by
the server.

A data refresh occurs on the data refresh server in the openverse-api project.
This is a task which imports data from the upstream Catalog database into the
API, copies contents to a new Elasticsearch index, and makes the index "live".
This process is necessary to make new content added to the Catalog by our
provider DAGs available on the frontend. You can read more in the [README](
https://github.com/WordPress/openverse-api/blob/main/ingestion_server/README.md
)

The TaskGroup generated by this factory allows us to schedule those refreshes through
Airflow. Since no two refreshes can run simultaneously, all tasks are initially
funneled through a special `data_refresh` pool with a single worker slot. To ensure
that tasks run in an acceptable order (ie the trigger step for one DAG cannot run if a
previously triggered refresh is still running), each DAG has the following
steps:

1. The `wait_for_data_refresh` step uses a custom Sensor that will wait until
none of the `external_dag_ids` (corresponding to the other data refresh DAGs)
are 'running'. A DAG is considered to be 'running' if it is itself in the
RUNNING state __and__ its own `wait_for_data_refresh` step has completed
successfully. The Sensor suspends itself and frees up the worker slot if
another data refresh DAG is running.

2. The `trigger_data_refresh` step then triggers the data refresh by POSTing
to the `/task` endpoint on the data refresh server with relevant data. A
successful response will include the `status_check` url used to check on the
status of the refresh, which is passed on to the next task via XCom.

3. Finally the `wait_for_data_refresh` task waits for the data refresh to be
complete by polling the `status_url`. Note this task does not need to be
able to suspend itself and free the worker slot, because we want to lock the
entire pool on waiting for a particular data refresh to run.

You can find more background information on this process in the following
issues and related PRs:

- [[Feature] Data refresh orchestration DAG](
https://github.com/WordPress/openverse-catalog/issues/353)
"""
import json
import logging
import os
import uuid
from collections.abc import Sequence
from urllib.parse import urlparse

from airflow.exceptions import AirflowException
from airflow.models.baseoperator import chain
from airflow.operators.python import PythonOperator
from airflow.providers.http.operators.http import SimpleHttpOperator
from airflow.providers.http.sensors.http import HttpSensor
from airflow.utils.task_group import TaskGroup
from requests import Response

from common.constants import XCOM_PULL_TEMPLATE
from common.sensors.single_run_external_dags_sensor import SingleRunExternalDAGsSensor
from data_refresh.data_refresh_types import DataRefresh


logger = logging.getLogger(__name__)


DATA_REFRESH_POOL = "data_refresh"


def response_filter_stat(response: Response) -> str:
    """
    Handle the response for the `get_current_index` task.

    This is used to extract the name of the current index that the concerned alias
    points to. This index name will be available via XCom in the downstream tasks.
    """
    index_name = response.json()["alt_names"]
    # Indices are named as '<media type>-<suffix>', so everything after the first
    # hyphen '-' is the suffix.
    _, index_suffix = index_name.split("-", maxsplit=1)
    return index_suffix


def response_filter_data_refresh(response: Response) -> str:
    """
    Handle the response for `trigger_data_refresh` task.

    This is used to grab the endpoint needed to poll for the status of the triggered
    data refresh. This information will then be available via XCom in the downstream
    tasks.
    """
    status_check_url = response.json()["status_check"]
    return urlparse(status_check_url).path


def response_check_wait_for_completion(response: Response) -> bool:
    """
    Handle the response for `wait_for_completion` Sensor.

    Processes the response to determine whether the task can complete.
    """
    data = response.json()

    if data["active"]:
        # The data refresh is still running. Poll again later.
        return False

    if data["error"]:
        raise AirflowException(
            "Ingestion server encountered an error during data refresh."
        )

    logger.info(f"Data refresh done with {data['progress']}% completed.")
    return True


def create_data_refresh_task_group(
    data_refresh: DataRefresh, external_dag_ids: Sequence[str]
):
    """
    Create the data refresh tasks.

    This factory method instantiates a DAG that will run the data refresh for
    the given `media_type`.

    A data refresh runs for a given media type in the API DB. It refreshes popularity
    data for that type, imports the data from the upstream DB in the Catalog, reindexes
    the data, and updates and reindex Elasticsearch.

    A data refresh can only be performed for one media type at a time, so the DAG
    must also use a Sensor to make sure that no two data refresh tasks run
    concurrently.

    It is intended that the data_refresh tasks, or at least the initial
    `wait_for_data_refresh` tasks, should be run in a custom pool with 1 worker
    slot. This enforces that no two `wait_for_data_refresh` tasks can start
    concurrently and enter a race condition.

    Required Arguments:

    data_refresh:     dataclass containing configuration information for the
                      DAG
    external_dag_ids: list of ids of the other data refresh DAGs. This DAG
                      will not run concurrently with any dependent DAG.
    """

    poke_interval = int(os.getenv("DATA_REFRESH_POKE_INTERVAL", 60 * 15))
    target_alias = data_refresh.media_type  # TODO: Change when using versioned aliases

    with TaskGroup(group_id="data_refresh") as data_refresh_group:
        tasks = []
        # Wait to ensure that no other Data Refresh DAGs are running.
        wait_for_data_refresh = SingleRunExternalDAGsSensor(
            task_id="wait_for_data_refresh",
            external_dag_ids=external_dag_ids,
            check_existence=True,
            poke_interval=poke_interval,
            mode="reschedule",
            pool=DATA_REFRESH_POOL,
        )
        tasks.append(wait_for_data_refresh)

        # Get the index currently mapped to our target alias, to delete later.
        get_current_index = SimpleHttpOperator(
            task_id="get_current_index",
            http_conn_id="data_refresh",
            endpoint=f"stat/{target_alias}",
            method="GET",
            response_check=lambda response: response.status_code == 200,
            response_filter=response_filter_stat,
        )
        tasks.append(get_current_index)

        # Generate a UUID suffix that will be used by the newly created index.
        generate_index_suffix = PythonOperator(
            task_id="generate_index_suffix", python_callable=lambda: uuid.uuid4().hex
        )
        tasks.append(generate_index_suffix)

        action_data_map: dict[str, dict] = {
            "ingest_upstream": {},
            "promote": {"alias": target_alias},
        }
        for action, action_post_data in action_data_map.items():
            with TaskGroup(group_id=action) as task_group:
                trigger = SimpleHttpOperator(
                    task_id=f"trigger_{action}",
                    http_conn_id="data_refresh",
                    endpoint="task",
                    method="POST",
                    headers={"Content-Type": "application/json"},
                    data=json.dumps(
                        action_post_data
                        | {
                            "model": data_refresh.media_type,
                            "action": action.upper(),
                            "index_suffix": XCOM_PULL_TEMPLATE.format(
                                generate_index_suffix.task_id, "return_value"
                            ),
                        }
                    ),
                    response_check=lambda response: response.status_code == 202,
                    response_filter=response_filter_data_refresh,
                )
                waiter = HttpSensor(
                    task_id=f"wait_for_{action}",
                    http_conn_id="data_refresh",
                    endpoint=XCOM_PULL_TEMPLATE.format(trigger.task_id, "return_value"),
                    method="GET",
                    response_check=response_check_wait_for_completion,
                    mode="reschedule",
                    poke_interval=poke_interval,
                    timeout=data_refresh.data_refresh_timeout.total_seconds(),
                )
                trigger >> waiter

            tasks.append(task_group)

        # Delete the alias' previous target index, now unused.
        delete_old_index = SimpleHttpOperator(
            task_id="delete_old_index",
            http_conn_id="data_refresh",
            endpoint="task",
            method="POST",
            headers={"Content-Type": "application/json"},
            data=json.dumps(
                {
                    "model": data_refresh.media_type,
                    "action": "DELETE_INDEX",
                    "index_suffix": XCOM_PULL_TEMPLATE.format(
                        get_current_index.task_id, "return_value"
                    ),
                }
            ),
            response_check=lambda response: response.status_code == 202,
        )
        tasks.append(delete_old_index)

        # ``tasks`` contains the following tasks and task groups:
        # wait_for_data_refresh
        # └─ get_current_index
        #    └─ ingest_upstream (trigger_ingest_upstream + wait_for_ingest_upstream)
        #       └─ promote (trigger_promote + wait_for_promote)
        #          └─ delete_old_index
        chain(*tasks)

    return data_refresh_group
