"""DB utils."""

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, MarketRegionTypeEnum
from apollo_main_db.spotify.models import (SpotifyAnalyticsAccountStreamInfo, SpotifyHotHitsPlaylistTrackStats,
                                           SpotifyMarketRank, SpotifyPlaylistTrackListHistory2,
                                           SpotifyPlaylistTrackListHistoryDates2, SpotifyTrack2, SpotifyViewPlaylist,
                                           ViewPlaylistTypeEnum)
from sqlalchemy import and_, bindparam, create_engine, func, or_, update
from sqlalchemy.orm import 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))


@utils.handle_errors()
def get_hot_hits_playlists() -> list:
    """Get hot hits playlist IDs and latest dates.

    """
    with session_scope() as session:
        return (
            session.query(SpotifyViewPlaylist.playlist_id, SpotifyViewPlaylist.last_date)
            .filter(SpotifyViewPlaylist.type == ViewPlaylistTypeEnum.HH)
            .all()
        )


@utils.single_column
@utils.handle_errors()
def check_playlists_in_history(playlist_ids: List[str]) -> list:
    """Check playlist IDs in history.

    """
    with session_scope() as session:
        return (
            session.query(SpotifyPlaylistTrackListHistoryDates2.playlist_id)
            .distinct()
            .filter(SpotifyPlaylistTrackListHistoryDates2.playlist_id.in_(playlist_ids))
            .all()
        )


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):
    """Handle S3 upload errors by removing non finished chunk files.

    """
    _, s3_path = args
    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, playlist_ids: List[str], hot_hits_date: date, s3_path: str):
    """Upload data chunk to S3.

    Args:
        table_type: Upload table type.
        playlist_ids: Playlist ID list.
        hot_hits_date: Latest hot hits date.
        s3_path (str): S3 upload path.

    """
    return exec_command(
        sql_consts.SQL.UPLOAD_S3,
        arguments={"id_list": tuple(playlist_ids), "date": hot_hits_date},
        replacements=dict(s3_path=f"{s3_path}h", columns=sql_consts.SQL.UPLOAD_S3_COLUMNS[table_type]),
    )


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,
        replacements=dict(
            s3_path=s3_path,
            data_table=sql_consts.Table.MAPPING[table_type],
            columns=sql_consts.SQL.DOWNLOAD_S3_COLUMNS[table_type],
        ),
    )


class AggregatedBase:
    table = None
    name = None

    @utils.single_column
    @utils.handle_errors()
    def get_playlists(self) -> list:
        """Get aggregated playlist IDs.

        """
        with session_scope() as session:
            return session.query(self.table.playlist_id).distinct().all()

    @utils.handle_errors()
    def delete_playlists(self, playlist_ids: List[str]) -> int:
        """Drop non hot hits playlists stats data.

        Args:
            playlist_ids: Playlist ID list to delete.

        Returns:
            int: Affected rows count.

        """
        with session_scope() as session:
            return (
                session.query(self.table)
                .filter(self.table.playlist_id.in_(playlist_ids))
                .delete(synchronize_session=False)
            )


def get_last_streams_date() -> date:
    """Get max streams table date.

    Returns:
        Last available date.

    """
    with session_scope() as session:
        return (
            session.query(func.max(SpotifyAnalyticsAccountStreamInfo.date))
            .filter(SpotifyAnalyticsAccountStreamInfo.account == 1)
            .first()
        )[0]


def get_markets_order_by_rank(date_from: date, date_to: date) -> List[str]:
    """Get markets ordered by rank.

    Args:
        date_from: Date from.
        date_to: Date to.

    Returns:
        Ordered list of markets.

    """
    with session_scope() as session:
        return [
            i[0]
            for i in (
                session.query(SpotifyAnalyticsAccountStreamInfo.market)
                .filter(SpotifyAnalyticsAccountStreamInfo.account == 1)
                .filter(SpotifyAnalyticsAccountStreamInfo.date >= date_from)
                .filter(SpotifyAnalyticsAccountStreamInfo.date <= date_to)
                .group_by(SpotifyAnalyticsAccountStreamInfo.market)
                .order_by(func.sum(SpotifyAnalyticsAccountStreamInfo.total_streams).desc())
            )
        ]


def get_playlists_markets():
    """Get HH playlists markets.

    Returns:
        HH playlists markets.

    """
    with session_scope() as session:
        return (
            session.query(SpotifyViewPlaylist.market_code)
            .filter(SpotifyViewPlaylist.type == ViewPlaylistTypeEnum.HH)
            .all()
        )


def set_playlists_data(playlists_mapping: Dict[str, Any], column_name: str):
    """Set HH playlists column values.

    Args:
        playlists_mapping: Playlist ID to column value mapping.
        column_name: Column name.

    """
    with session_scope() as session:
        query = (
            update(SpotifyViewPlaylist)
            .where(SpotifyViewPlaylist.playlist_id == bindparam("playlist_id"))
            .where(SpotifyViewPlaylist.type == ViewPlaylistTypeEnum.HH)
            .values({getattr(SpotifyViewPlaylist, column_name): bindparam(column_name)})
        )
        session.execute(
            query,
            [{"playlist_id": playlist_id, column_name: value} for playlist_id, value in playlists_mapping.items()],
        )


def set_playlists_ranks(market_rank_mapping: Dict[str, int], column_name: str = "rank"):
    """Set HH markets ranks.

    Args:
        market_rank_mapping: Market code to rank mapping.
        column_name: Column name to update.

    """
    with session_scope() as session:
        query = (
            update(SpotifyMarketRank)
            .where(SpotifyMarketRank.market_code == bindparam("market_code"))
            .where(SpotifyMarketRank.type == MarketRegionTypeEnum.HH)
            .values({getattr(SpotifyMarketRank, column_name): bindparam(column_name)})
        )
        session.execute(
            query,
            [{"market_code": market_code, column_name: value} for market_code, value in market_rank_mapping.items()],
        )


def set_playlists_dates(playlists_dates: Dict[str, date]):
    """Set HH playlists last dates.

    Args:
        playlists_dates: Playlist ID to date mapping.

    """
    set_playlists_data(playlists_dates, "last_date")


class Stats(AggregatedBase):
    table = SpotifyHotHitsPlaylistTrackStats
    name = consts.Table.STATS

    @utils.handle_errors()
    def get_previous_positions(
        self,
        playlist_id: str,
        position_date_from: date,
        position_date_to: date,
        latest_date_from: Optional[date] = None,
        latest_date_to: Optional[date] = None,
        isrc_list: Optional[List[str]] = None,
    ):
        """Get previous positions for tracks in HH playlist.

        Args:
            playlist_id: HH playlist ID.
            position_date_from: Previous position date from.
            position_date_to: Previous position date to.
            latest_date_from: Latest position date interval from.
            latest_date_to: Latest position date interval to.
            isrc_list: ISRC list.

        Returns:
            Previous positions and dates per ISRC.

        """
        with session_scope() as session:
            select_columns = [
                SpotifyTrack2.isrc,
                func.min(SpotifyPlaylistTrackListHistory2.playlist_index).label("playlist_index"),
                SpotifyPlaylistTrackListHistory2.date,
            ]

            if latest_date_from and latest_date_to:
                select_columns.append(self.table.latest_position)

            query = (
                session.query(*select_columns)
                .select_from(SpotifyPlaylistTrackListHistory2)
                .join(SpotifyTrack2, SpotifyTrack2.id == SpotifyPlaylistTrackListHistory2.track_id)
                .filter(SpotifyPlaylistTrackListHistory2.playlist_id == playlist_id)
                .filter(SpotifyPlaylistTrackListHistory2.date.between(position_date_from, position_date_to))
            )

            if isrc_list is not None:
                query = query.filter(SpotifyTrack2.isrc.in_(isrc_list))

            if latest_date_from and latest_date_to:
                query = query.join(
                    self.table,
                    and_(
                        self.table.isrc == SpotifyTrack2.isrc,
                        self.table.playlist_id == SpotifyPlaylistTrackListHistory2.playlist_id,
                        self.table.latest_date > SpotifyPlaylistTrackListHistory2.date,
                        self.table.latest_date.between(latest_date_from, latest_date_to),
                    ),
                )

            return query.group_by(SpotifyTrack2.isrc, SpotifyPlaylistTrackListHistory2.date).all()

    @utils.handle_errors()
    def get(self, playlist_id: str, isrc_list: List[str], current_date: date) -> List[SpotifyHotHitsPlaylistTrackStats]:
        """Get stats records that may have changes.

        Args:
            playlist_id: Playlist ID.
            isrc_list: Track ISRC list.
            current_date: Current processing date.

        Returns:
            Stats records.

        """
        with session_scope() as session:
            result = (
                session.query(self.table)
                .filter(self.table.playlist_id == playlist_id)
                .filter(or_(self.table.isrc.in_(isrc_list), self.table.latest_date >= current_date))
                .all()
            )
            session.expunge_all()
            return result

    @utils.handle_errors()
    def add(self, playlist_id: str, stats_list: List[dict]):
        """Create a new hot hits stats record.

        Args:
            playlist_id: Playlist ID.
            stats_list: Stats data list.

        """
        records = []
        for item in stats_list:
            record = SpotifyHotHitsPlaylistTrackStats()
            record.playlist_id = playlist_id
            record.isrc = item["isrc"]
            record.entry_date = item["entry_date"]
            record.peak_date = item["peak_date"]
            record.peak_position = item["peak_position"]
            record.previous_date = item["previous_date"]
            record.previous_position = item["previous_position"]
            record.latest_date = item["latest_date"]
            record.latest_position = item["latest_position"]
            record.added_date = item["added_date"]
            record.in_history = item["in_history"]
            records.append(record)

        with session_scope() as session:
            return session.bulk_save_objects(records)

    @utils.handle_errors()
    def update(
        self,
        playlist_id: str,
        stats_data: Union[List[dict], Dict[str, List[dict]]],
        fields: Optional[List[str]] = None,
    ):
        """Set HH stats records values.

        Args:
            playlist_id: Playlist ID.
            stats_data: Stats values.
            fields: Fields names to updated.

        """
        if isinstance(stats_data, dict):
            stats_data = list(stats_data.values())
        if not fields:
            fields = [i for i in stats_data[0].keys() if i != "isrc"]
        with session_scope() as session:
            query = (
                update(self.table)
                .where(self.table.playlist_id == playlist_id)
                .where(self.table.isrc == bindparam("isrc"))
                .values({getattr(self.table, field): bindparam(field) for field in fields})
            )
            return session.execute(query, stats_data)

    @utils.handle_errors()
    def delete(self, playlist_id: str, isrc_list: List[str]) -> int:
        """Delete stats records.

        Args:
            playlist_id: Playlist ID.
            isrc_list: ISRC list.

        Returns:
            Removed records count.

        """
        with session_scope() as session:
            query = (
                session.query(self.table)
                .filter(self.table.playlist_id == playlist_id)
                .filter(self.table.isrc.in_(isrc_list))
            )
            return query.delete(synchronize_session=False)

    @utils.handle_errors()
    def delete_extra(self, playlist_id: str, isrc_list: List[str], current_date: date) -> int:
        """Delete stats records.

        Args:
            playlist_id: Playlist ID.
            isrc_list: Current track ISRC list.
            current_date: Current processing date.

        Returns:
            Removed records count.

        """
        with session_scope() as session:
            query = (
                session.query(self.table)
                .filter(self.table.playlist_id == playlist_id)
                .filter(self.table.isrc.notin_(isrc_list))
                .filter(self.table.entry_date == current_date)
                .filter(self.table.latest_date >= current_date)
            )
            return query.delete(synchronize_session=False)


stats = Stats()

db = {consts.Table.STATS: stats}