import json

from airflow.providers.http.operators.http import SimpleHttpOperator
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.utils.task_group import TaskGroup

from common.auth.atlas import AtlasAuth
from common.operators.snowflake import XCOMSnowflakeOperator


__all__ = ["PreparationTaskManager"]


class PreparationTaskManager:

    def get_task(self):
        with TaskGroup("get_excluded") as get_excluded:
            SimpleHttpOperator(
                task_id="partial_match_labels",
                method="GET",
                endpoint="label_list",
                data={"full_match": False},
                http_conn_id="cms_http",
                response_filter=lambda response: "|".join(json.loads(response.text)),
                auth_type=AtlasAuth
            )
            SimpleHttpOperator(
                task_id="full_match_labels",
                method="GET",
                endpoint="label_list",
                data={"full_match": True},
                http_conn_id="cms_http",
                response_filter=lambda response: json.loads(response.text),
                auth_type=AtlasAuth
            )
            SimpleHttpOperator(
                task_id="artists",
                method="GET",
                endpoint="artist_list",
                http_conn_id="cms_http",
                response_filter=lambda response: "|".join(json.loads(response.text)),
                auth_type=AtlasAuth
            )
        prepare_nsa = XCOMSnowflakeOperator(
            task_id="non_signed_artists",
            sql="preparation/artist/prepare_nsa.sql",
            params={
                "table_name": "T_NON_SIGNED_ARTIST",
                "table_name_alt": "T_NON_SIGNED_ARTIST_ALT",
                "diff_table_name": "DIFF_NOT_SIGNED_ARTIST"
            },
            xcom_parameters={
                "EXCLUDED_PARTIAL_MATCH_LABELS": {"task_ids": "preparation.get_excluded.partial_match_labels"},
                "EXCLUDED_FULL_MATCH_LABELS": {"task_ids": "preparation.get_excluded.full_match_labels"},
                "EXCLUDED_ARTISTS": {"task_ids": "preparation.get_excluded.artists"},
            }
        )

        prepare_nsa_genres = SnowflakeOperator(
            task_id="non_signed_artists_genres",
            sql="preparation/artist/prepare_genre.sql",
            params={
                "table_name": "NON_SIGNED_ARTIST_GENRE",
                "table_name_alt": "NON_SIGNED_ARTIST_GENRE_ALT"
            }
        )

        prepare_nst = SnowflakeOperator(
            task_id="non_signed_tracks",
            sql="preparation/track/prepare_nst.sql",
            params={
                "table_name": "T_NON_SIGNED_TRACK",
                "table_name_alt": "T_NON_SIGNED_TRACK_ALT"
            }
        )

        prepare_nst_genres = SnowflakeOperator(
            task_id="non_signed_tracks_genres",
            sql="preparation/track/prepare_genre.sql",
            params={
                "table_name": "NON_SIGNED_TRACK_GENRE",
                "table_name_alt": "NON_SIGNED_TRACK_GENRE_ALT"
            }
        )

        get_excluded >> prepare_nsa >> prepare_nsa_genres
        prepare_nsa >> prepare_nst >> prepare_nst_genres

        return prepare_nsa_genres, prepare_nst_genres