import json

import pendulum
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


def non_signed_data_handler(cursor) -> str:
    result = {
        "artist": {},
        "track": {}
    }

    if cursor.description is not None:
        for key, value in cursor.fetchone().items():
            key = key.lower()

            if "nsa" in key:
                group = "artist"
            elif "nst" in key:
                group = "track"
            else:
                group = None

            if group is not None:
                result[group][key] = value

    return json.dumps(result)


def dropped_artist_data_handler(cursor) -> str:
    result = []
    if cursor.description is not None:
        for row in cursor.fetchall():
            for key, value in row.items():
                result.append({key.lower(): value})

    return json.dumps(result)


def dropped_artist_cms_data_filter(response):
    created_at = json.loads(response.text).get("created_at")

    if not created_at:
        created_at = (pendulum.today(tz=pendulum.UTC) - pendulum.duration(days=1)).isoformat()

    return created_at


class CMSReportTaskManager:
    def get_task(self):
        with TaskGroup("cms_data") as cms_data:
            with TaskGroup("non_signed"):
                get_latest_non_signed_cms_data = SimpleHttpOperator(
                    task_id="get_from_cms",
                    method="GET",
                    endpoint="last_data",
                    http_conn_id="cms_http",
                    response_filter=lambda response: json.loads(response.text),
                    auth_type=AtlasAuth
                )

                get_non_signed_data = SnowflakeOperator(
                    task_id="get_from_db",
                    sql="cms_report/get_non_signed_data.sql",
                    handler=non_signed_data_handler
                )

                upload_latest_non_signed_cms_data = SimpleHttpOperator(
                    task_id="upload_to_cms",
                    method="POST",
                    endpoint="last_data",
                    data="{{ ti.xcom_pull(task_ids='cms_data.non_signed.get_from_db') }}",
                    headers={
                        "Content-Type": "application/json"
                    },
                    http_conn_id="cms_http",
                    auth_type=AtlasAuth
                )

                get_latest_non_signed_cms_data >> get_non_signed_data >> upload_latest_non_signed_cms_data

            # TODO: Figure out how to work with large result
            # with TaskGroup("dropped_artist"):
            #     get_latest_dropped_artist_cms_data = SimpleHttpOperator(
            #         task_id="get_from_cms",
            #         method="GET",
            #         endpoint="weekly_report_data",
            #         http_conn_id="cms_http",
            #         response_filter=dropped_artist_cms_data_filter,
            #         auth_type=AtlasAuth
            #     )
            #
            #     get_dropped_artist_data = SnowflakeOperator(
            #         task_id="get_from_db",
            #         sql="cms_report/get_dropped_artist_data.sql",
            #         handler=dropped_artist_data_handler
            #     )
            #
            #     upload_latest_dropped_artist_cms_data = SimpleHttpOperator(
            #         task_id="upload_to_cms",
            #         method="POST",
            #         endpoint="weekly_report_data",
            #         data="{{ ti.xcom_pull(task_ids='cms_data.dropped_artist.get_from_db') }}",
            #         headers={
            #             "Content-Type": "application/json"
            #         },
            #         http_conn_id="cms_http",
            #         auth_type=AtlasAuth
            #     )
            #
            #     get_latest_dropped_artist_cms_data >> get_dropped_artist_data >> upload_latest_dropped_artist_cms_data

        return cms_data