import uuid

from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from airflow.utils.context import Context
from airflow.utils.state import State


__all__ = ["notify_dag_success", "notify_dag_failed"]


# TODO: in airflow version 2.6.0+ change to airflow.providers.slack.notifications.slack_notifier.send_slack_notification

def notify_dag_failed(context: Context):
    tis_dagrun = context["ti"].get_dagrun().get_task_instances()
    failed_tasks = []
    for ti in tis_dagrun:
        if ti.state == State.FAILED:
            # Adding log url
            failed_tasks.append(f"<{ti.log_url}|{ti.task_id}>")

    dag = context["task_instance"].dag_id
    exec_date = context.get('execution_date')

    blocks = [
        {
            "type": "section",
            "text": {
                "type": "mrkdwn",
                "text": ":octagonal_sign: ETL process failed. @channel"
            }
        },
        {
            "type": "section",
            "block_id": f"section_{uuid.uuid4()}",
            "text": {
                "type": "mrkdwn",
                "text": f"*Dag*: {dag} \n *Execution Time*: {exec_date}"
            }
        },
        {
            "type": "section",
            "text": {
                "type": "mrkdwn",
                "text": f"Failed Tasks: {', '.join(failed_tasks)}"
            }
        }
    ]
    notifier = SlackAPIPostOperator(
        task_id="slack_notify",
        slack_conn_id="slack_api_default",
        channel="#airflow-etl",
        blocks=blocks,
        text="ETL process failed. @channel"
    )
    notifier.execute()


def notify_dag_success(context: Context):
    dag = context["task_instance"].dag_id
    exec_date = context.get('execution_date')
    blocks = [
        {
            "type": "section",
            "text": {
                "type": "mrkdwn",
                "text": ":white_check_mark: ETL processing finished successfully."
            }
        },
        {
            "type": "section",
            "block_id": f"section_{uuid.uuid4()}",
            "text": {
                "type": "mrkdwn",
                "text": f"*Dag*: {dag} \n *Execution Time*: {exec_date}"
            }
        }
    ]
    notifier = SlackAPIPostOperator(
        task_id="slack_notify",
        slack_conn_id="slack_api_default",
        channel="#airflow-etl",
        blocks=blocks,
        text="ETL processing finished successfully."
    )
    notifier.execute()