from typing import Any, Dict

from requests import Session  # type: ignore

from config import CMS_API_CONFIG
from steps.base import BaseStep
from steps.mixins.snowflake import SnowflakeMixin
from utils.external_clients.cms.client import CMSClient
from utils.external_clients.cms.client import Config as CMSConfig
from utils.failsafe import failsafe

from ..utils.cleanup import Cleanup


__all__ = ["DataUpload"]


class DataUpload(SnowflakeMixin, BaseStep):
    depends_on = {Cleanup}
    non_signed_artist_tbl = "T_NON_SIGNED_ARTIST"
    diff_not_signed_artist_tbl = "DIFF_NOT_SIGNED_ARTIST"
    non_signed_track_tbl = "T_NON_SIGNED_TRACK"
    total_nsa_key = "total_nsa"
    dropped_nsa_key = "dropped_nsa"
    new_nsa_key = "new_nsa"
    total_nst_key = "total_nst"
    new_nst_key = "new_nst"

    def _get_total_artist_count(self) -> int:
        query = f"SELECT COUNT(*) AS {self.total_nsa_key} FROM DNA.DNA_PUBLIC.{self.non_signed_artist_tbl}"

        result = self.run_raw_query(query, is_async=False)
        result = result[0].get(self.total_nsa_key)
        return result

    def _get_dropped_off_artist(self, created_at: str) -> int:
        query = f"""
            SELECT COUNT(*) AS {self.dropped_nsa_key} FROM DNA.DNA_PUBLIC.{self.diff_not_signed_artist_tbl}
            WHERE created_at > '{created_at}'
        """

        result = self.run_raw_query(query, is_async=False)
        result = result[0].get(self.dropped_nsa_key)
        return result

    def _get_total_track_count(self) -> int:
        query = f"SELECT COUNT(*) AS {self.total_nst_key} FROM DNA.DNA_PUBLIC.{self.non_signed_track_tbl}"

        result = self.run_raw_query(query, is_async=False)
        result = result[0].get(self.total_nst_key)
        return result

    def _pre_run(self):
        super()._pre_run()
        session = Session()
        self.cms_api = CMSClient(session=session, config=CMSConfig(**CMS_API_CONFIG))

    @failsafe
    def process(self) -> Dict[str, Any]:
        last_data = self.cms_api.get_last_data()

        artist_data = {}
        total_nsa = self._get_total_artist_count()
        artist_data[self.total_nsa_key] = total_nsa
        if last_data["artist"]:
            artist_data[self.dropped_nsa_key] = self._get_dropped_off_artist(last_data["artist"]["created_at"])
            artist_data[self.new_nsa_key] = total_nsa - last_data["artist"]["total_nsa"]

        track_data = {}
        total_nst = self._get_total_track_count()
        track_data[self.total_nst_key] = total_nst
        if last_data["track"]:
            track_data[self.new_nst_key] = total_nst - last_data["track"]["total_nst"]

        post_last_data = {"artist": artist_data, "track": track_data}

        return self.cms_api.post_last_data(post_last_data)