"""DB queries."""

from contextlib import contextmanager
from datetime import date, datetime
from typing import Any, Dict, List, Optional, Union

from apollo_main_db.apollo.models import ApolloKeyValueStorage, BuzzUser
from apollo_main_db.spotify.models import (SpotifyPlaylist, SpotifyPlaylistDigestFollowersPlaylist,
                                           SpotifyPlaylistDigestStreamsPlaylist, SpotifyPlaylistFollowers,
                                           SpotifyPlaylistFollowersHistory, SpotifyPlaylistFollowersWeekHistory,
                                           SpotifyPlaylistStreamSummaryGlobal, SpotifyPlaylistStreamsWeekHistory)
from sqlalchemy import and_, bindparam, case, create_engine, func, literal, or_, update
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.orm import Query, Session, scoped_session, sessionmaker

import config
import s3
import utils
from constants import common as consts
from constants import sql as sql_consts
from logger import logger

engine = create_engine(
    "mysql+pymysql://{user}:{password}@{host}:{port}/{database_name}".format(
        user=config.MySQL.USER,
        password=config.MySQL.PASSWORD,
        host=config.MySQL.HOST,
        port=config.MySQL.PORT,
        database_name=config.MySQL.DATABASE,
    ),
    pool_recycle=config.MySQL.POOL_RECYCLE,
    pool_size=config.MySQL.POOL_SIZE,
)
_session_factory = scoped_session(sessionmaker(bind=engine))
_session = _session_factory()


@contextmanager
def session_scope():
    """Provide a transactional scope around a series of operations."""
    try:
        yield _session
        _session.commit()
    except Exception:
        _session.rollback()
        raise
    finally:
        _session.close()


def get_value(key: str) -> str or None:
    """Get value by key from key-value storage.

    Args:
        key: Key.

    Returns:
        Value.

    """
    with session_scope() as session:
        result = session.query(ApolloKeyValueStorage.value).filter(ApolloKeyValueStorage.key == key).first()
        return result[0] if result else None


def get_date_value(key: str, default_value: date or None = None) -> date or None:
    """Get value from DB Apollo key-value storage.

    Args:
        key: Key.
        default_value: Default value if None.

    Returns:
        Value.

    """
    value = get_value(key)
    if not value:
        return default_value
    return datetime.strptime(value, consts.DEFAULT_DATE_FORMAT).date()


def set_value(key: str, value: str):
    """Set value for a key to key-value storage.

    Args:
        key: Key.
        value: Value.

    """
    with session_scope() as session:
        (
            session.query(ApolloKeyValueStorage)
            .filter(ApolloKeyValueStorage.key == key)
            .update({ApolloKeyValueStorage.value: value}, synchronize_session=False)
        )


def set_date_value(key: str, value: date):
    """Set value to DB Apollo key-value storage.

    Args:
        key: Key.
        value: Value.

    """
    set_value(key, value.strftime(consts.DEFAULT_DATE_FORMAT))


def exec_command(sql_text: str, arguments: Dict[str, Any] or None = None, replacements: dict or None = None) -> int:
    """Execute DB change command.

    Args:
        sql_text: SQL command.
        arguments: Command parameters.
        replacements: SQL command substitutions.

    Return:
        int: Row affected count.

    """
    if replacements:
        sql_text = sql_text.format(**replacements)
    with session_scope() as session:
        row_count = session.execute(sql_text, arguments)
    return row_count if isinstance(row_count, int) else row_count.rowcount


@utils.handle_errors()
def exec_command_retry(
    sql_text: str, arguments: Dict[str, Any] or None = None, replacements: dict or None = None
) -> int:
    """Execute DB change command.

    Args:
        sql_text: SQL command.
        arguments: Command parameters.
        replacements: SQL command substitutions.

    Return:
        int: Row affected count.

    """
    return exec_command(sql_text, arguments, replacements)


def handle_s3_non_finished_chunk(*args, **kwargs):
    """Handle S3 upload errors by removing non finished chunk files."""
    s3_path = args[1]
    if len(args) > 2:
        file_prefix = args[2]
        s3_path = f"{s3_path}{file_prefix}"
    logger.debug(f"Removing files in {s3_path}")
    s3.delete_folder_files(s3.get_relative_by_full(s3_path))


@utils.handle_errors(custom_handler=handle_s3_non_finished_chunk)
def upload_to_s3(table_type: str, s3_path: str, file_prefix: str = "f", **kwargs) -> int:
    """Upload data chunk to S3.

    Args:
        table_type: Upload table type.
        s3_path: S3 upload path.
        file_prefix: S3 file name prefix.

    Results:
        Row count.

    """
    return exec_command(
        sql_consts.SQL.UPLOAD_S3[table_type], arguments=kwargs, replacements=dict(s3_path=f"{s3_path}{file_prefix}")
    )


def download_from_s3(table_type, s3_path: str) -> int:
    """Download data chunks from S3 to MySQL.

    Args:
        table_type: Download table type.
        s3_path (str): S3 download path.

    Returns:
        int: Rows count.

    """
    return exec_command_retry(sql_consts.SQL.DOWNLOAD_S3[table_type], replacements=dict(s3_path=s3_path))


class AggregatedBase:
    table: Union[SpotifyPlaylistStreamsWeekHistory, SpotifyPlaylistFollowersWeekHistory] = None
    playlists_table: Union[SpotifyPlaylistDigestStreamsPlaylist, SpotifyPlaylistDigestFollowersPlaylist] = None
    metrics_table: Union[SpotifyPlaylistStreamSummaryGlobal, SpotifyPlaylistFollowers] = None
    metrics_column_name: str = None
    name: str = None
    key_name: str = None

    def _get_top_playlists_query(
        self,
        session: Session,
        metrics_table: Union[
            SpotifyPlaylistStreamSummaryGlobal,
            SpotifyPlaylistFollowers,
            SpotifyPlaylistFollowersHistory,
        ] = None,
    ) -> Query:
        """Get top playlists query.

        Args:
            session: SQL session.
            metrics_table: Metric table.

        Returns:
            Top playlists query.

        """
        if not metrics_table:
            metrics_table = self.metrics_table
        metrics_column = getattr(metrics_table, self.metrics_column_name)
        return (
            session.query(metrics_table.playlist_id)
            .join(SpotifyPlaylist, SpotifyPlaylist.id == metrics_table.playlist_id)
            .outerjoin(BuzzUser, BuzzUser.user_name == SpotifyPlaylist.user_name)
            .filter(SpotifyPlaylist.removed == 0)
            .filter(
                and_(
                    or_(SpotifyPlaylist.buzz_category_id != 6, SpotifyPlaylist.buzz_category_id.is_(None)),
                    or_(BuzzUser.category_id != 6, BuzzUser.category_id.is_(None)),
                )
            )
            .filter(
                or_(
                    and_(
                        or_(SpotifyPlaylist.buzz_category_id != 5, SpotifyPlaylist.buzz_category_id.is_(None)),
                        or_(BuzzUser.category_id != 5, BuzzUser.category_id.is_(None)),
                    ),
                    SpotifyPlaylist.name.notlike("%radio"),
                )
            )
            .filter(metrics_table.playlist_id != "")
            .filter(
                or_(
                    metrics_column >= config.PLAYLISTS_NON_CATEGORIZED_MIN_METRIC_COUNT[self.name],
                    and_(
                        metrics_column >= config.PLAYLISTS_CATEGORIZED_MIN_METRIC_COUNT[self.name],
                        or_(SpotifyPlaylist.buzz_category_id.isnot(None), BuzzUser.category_id.isnot(None)),
                    )

                )
            )
        )

    @utils.single_column
    @utils.handle_errors()
    def get_aggregated_playlists(self, current_date: Optional[date] = None) -> List[str]:
        """Get aggregated playlist IDs.

        Args:
            current_date: Date to update.

        Returns:
            Playlist ID list.

        """
        with session_scope() as session:
            query = session.query(self.playlists_table.playlist_id)
            if current_date:
                query = query.filter(self.playlists_table.latest_date < current_date)
            return query.all()

    @utils.single_column
    @utils.handle_errors()
    def get_top_playlists(self) -> List[str]:
        """Get playlist IDs with more than some number of key metric.

        Returns:
            All available playlist ID list.

        """
        with session_scope() as session:
            return self._get_top_playlists_query(session).all()

    def set_latest_date(self, latest_date: date):
        """Set latest date.

        """
        set_date_value(self.key_name, latest_date)

    def upsert_data(self, data: List[dict]):
        """Save streams/followers for two weeks into weeks history.

        Args:
            data: List of streams/followers weeks history records.

        """
        with session_scope() as session:
            insert_stmt = insert(self.table).values(data)
            on_conflict_stmt = insert_stmt.on_duplicate_key_update(PlaylistId=insert_stmt.inserted.PlaylistId)
            session.execute(on_conflict_stmt)

    def save_data(self, data: List[SpotifyPlaylistFollowersWeekHistory or SpotifyPlaylistStreamsWeekHistory]):
        """Save streams/followers for two weeks into weeks history.

        Args:
            data: List of streams/followers weeks history records.

        """
        with session_scope() as session:
            session.bulk_save_objects(data)

    @utils.handle_errors()
    def insert_metric_playlists(self, missing_playlists: List[str], latest_date: date) -> int:
        """Insert new playlists.

        Args:
            missing_playlists: Playlists to add.
            latest_date: Latest date.

        Returns:
            Row count.

        """
        if not missing_playlists:
            return 0

        with session_scope() as session:
            query_data = (
                session.query(
                    SpotifyPlaylist.id.label("playlist_id"),
                    func.COALESCE(SpotifyPlaylist.buzz_category_id, BuzzUser.category_id, 0).label("category_id"),
                    case(
                        [
                            (
                                func.COALESCE(SpotifyPlaylist.country_code, BuzzUser.country_code, "").in_(
                                    ("", "_gl", "null")
                                ),
                                consts.WORLDWIDE,
                            )
                        ],
                        else_=func.IFNULL(SpotifyPlaylist.country_code, BuzzUser.country_code),
                    ).label("country_code"),
                    literal(latest_date).label("latest_date"),
                )
                .outerjoin(BuzzUser, BuzzUser.user_name == SpotifyPlaylist.user_name)
                .filter(SpotifyPlaylist.id.in_(missing_playlists))
                .subquery()
            )

            query_insert = insert(self.playlists_table).from_select(
                [
                    self.playlists_table.playlist_id,
                    self.playlists_table.category_id,
                    self.playlists_table.country_code,
                    self.playlists_table.latest_date,
                ],
                query_data,
            )
            query_on_conflict = query_insert.on_duplicate_key_update(LatestDate=query_insert.inserted.LatestDate)

            return session.execute(query_on_conflict).rowcount

    @utils.handle_errors()
    def get_metric_changed_playlists(self) -> List[dict]:
        """Get changed playlists.

        Returns:
            Changed playlists and new values.

        """
        with session_scope() as session:
            result = (
                session.query(
                    self.playlists_table.playlist_id.label("playlist_id"),
                    case(
                        [
                            (
                                func.COALESCE(SpotifyPlaylist.country_code, BuzzUser.country_code, "").in_(
                                    ("_gl", "", "null")
                                ),
                                consts.WORLDWIDE,
                            ),
                        ],
                        else_=func.IFNULL(SpotifyPlaylist.country_code, BuzzUser.country_code),
                    ).label("country_code"),
                    func.COALESCE(SpotifyPlaylist.buzz_category_id, BuzzUser.category_id, 0).label("category_id"),
                )
                .join(SpotifyPlaylist, SpotifyPlaylist.id == self.playlists_table.playlist_id)
                .outerjoin(BuzzUser, BuzzUser.user_name == SpotifyPlaylist.user_name)
                .filter(
                    or_(
                        case(
                            [
                                (
                                    func.COALESCE(SpotifyPlaylist.country_code, BuzzUser.country_code, "").in_(
                                        ("_gl", "", "null")
                                    ),
                                    consts.WORLDWIDE,
                                ),
                            ],
                            else_=func.IFNULL(SpotifyPlaylist.country_code, BuzzUser.country_code),
                        ) != self.playlists_table.country_code,
                        func.COALESCE(
                            SpotifyPlaylist.buzz_category_id, BuzzUser.category_id, 0
                        ) != self.playlists_table.category_id,
                    )
                )
                .all()
            )
        return [i._asdict() for i in result]

    @utils.handle_errors()
    def update_metric_playlists(self, changed_playlists: List[dict]) -> int:
        """Update changed playlists.

        Args:
            changed_playlists: Changed playlists and new values.

        """
        with session_scope() as session:
            query_update = (
                update(self.playlists_table)
                .values(CountryCode=bindparam("country_code"), CategoryId=bindparam("category_id"))
                .where(self.playlists_table.playlist_id == bindparam("playlist_id"))
            )

            return session.execute(query_update, changed_playlists).rowcount

    @utils.handle_errors()
    def set_latest_dates(self, data: List[dict]) -> int:
        """Update playlists latest dates.

        Args:
            data: Playlist ID to latest date mapping.

        """
        with session_scope() as session:
            query_update = (
                update(self.playlists_table)
                .values(LatestDate=bindparam("latest_date"))
                .where(self.playlists_table.playlist_id == bindparam("playlist_id"))
            )

            return session.execute(query_update, data).rowcount


class Streams(AggregatedBase):
    table = SpotifyPlaylistStreamsWeekHistory
    playlists_table = SpotifyPlaylistDigestStreamsPlaylist
    metrics_table = SpotifyPlaylistStreamSummaryGlobal
    metrics_column_name = "streams_7_days"
    name = consts.Table.STREAMS
    key_name = consts.ApolloKey.PLAYLIST_DIGEST_STREAMS_DATE


class Followers(AggregatedBase):
    table = SpotifyPlaylistFollowersWeekHistory
    playlists_table = SpotifyPlaylistDigestFollowersPlaylist
    metrics_table = SpotifyPlaylistFollowers
    metrics_column_name = "followers"
    name = consts.Table.FOLLOWERS
    key_name = consts.ApolloKey.PLAYLIST_DIGEST_FOLLOWERS_DATE

    @utils.handle_errors()
    @utils.as_date
    def get_latest_date(self):
        """Get the latest available date.

        Returns:
            The latest date.

        """
        with session_scope() as session:
            result = session.query(func.max(SpotifyPlaylistFollowersHistory.date)).first()
            return result[0]

    @utils.single_column
    @utils.handle_errors()
    def get_top_playlists(self) -> list:
        with session_scope() as session:
            if config.PLAYLIST_LIST_FOLLOWERS_DATE:
                return self._get_top_playlists_query(
                    session, metrics_table=SpotifyPlaylistFollowersHistory
                ).filter(SpotifyPlaylistFollowersHistory.date == config.PLAYLIST_LIST_FOLLOWERS_DATE).all()
            else:
                return self._get_top_playlists_query(session).all()

    @utils.handle_errors()
    def get_data(self, playlist_ids: List[str], current_date: date, day_7_date: date, day_14_date: date) -> list:
        """Get followers for chosen playlists and dates.

        Args:
            playlist_ids: Playlist ID list.
            current_date: Current date.
            day_7_date: Current - 7 days date.
            day_14_date: Current - 14 days date.

        Returns:
            Followers per playlist ID, date.

        """
        with session_scope() as session:
            return (
                session.query(
                    SpotifyPlaylistFollowersHistory.playlist_id,
                    SpotifyPlaylistFollowersHistory.date,
                    SpotifyPlaylistFollowersHistory.followers,
                )
                .filter(SpotifyPlaylistFollowersHistory.playlist_id.in_(playlist_ids))
                .filter(SpotifyPlaylistFollowersHistory.date.in_([current_date, day_7_date, day_14_date]))
                .all()
            )


streams = Streams()
followers = Followers()

db = {consts.Table.STREAMS: streams, consts.Table.FOLLOWERS: followers}