import logging
import os
import sys
import time
from datetime import datetime, timedelta
from typing import Generator

from apollo_main_db.apple import AppleMusicContainerStreams, AppleMusicContainerStreamSummary, ApplePlaylist, \
    AppleWeeklyTopPlaylist
from sqlalchemy import and_, func, or_

from config import session
from constants import DISTRIBUTOR_IDS, TOP_PLAY_LISTS_NUMBER


# set up logging
LOGGER_LEVEL = getattr(logging, os.environ.get('LOGGER_LEVEL', 'INFO'))
logger = logging.getLogger()
logger.setLevel(LOGGER_LEVEL)
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setLevel(logging.DEBUG)
logger.addHandler(stream_handler)


def date_range(start_date: datetime.date, end_date: datetime.date) -> Generator[datetime.date, None, None]:
    """Generator for dates between two dates

    Args:
        start_date (datetime.date): Start date.
        end_date (datetime.date): End date.

    Returns:
        Generator[datetime.date]: Dates.

    """
    for n in range(int((start_date - end_date).days) + 1):
        yield start_date - timedelta(n)


def back_fill_data(start_date: datetime.date, end_date: datetime.date):
    """Back fill data to AppleWeeklyTopPlaylist for date range

    Args:
        start_date (datetime.date): Start date of filling period.
        end_date (datetime.date): End date of filling period.

    Returns:
        None

    """
    logger.info(f'Script started for dates: {start_date} - {end_date}')
    country_codes = session.query(
        AppleMusicContainerStreamSummary.country_code).distinct(
        AppleMusicContainerStreamSummary.country_code)

    for week_start_date in date_range(start_date, end_date):
        week_end_date = week_start_date - timedelta(days=6)

        dates = [
            and_(AppleMusicContainerStreams.date == date)
            for date in date_range(week_start_date, week_end_date)]

        for country_code, in country_codes:
            start_time = time.time()
            session.query(AppleWeeklyTopPlaylist).filter(
                AppleWeeklyTopPlaylist.country_code == country_code,
                AppleWeeklyTopPlaylist.date == week_start_date,
            ).delete()

            query = session.query(
                func.sum(AppleMusicContainerStreams.streams).label('sum'),
                AppleMusicContainerStreams.container_id
            ).join(
                ApplePlaylist,
                ApplePlaylist.id == AppleMusicContainerStreams.container_id
            ).filter(
                AppleMusicContainerStreams.distributor_id.in_(DISTRIBUTOR_IDS),
                or_(*dates)
            ).group_by(
                AppleMusicContainerStreams.container_id
            ).order_by(func.sum(AppleMusicContainerStreams.streams).desc())

            if country_code != 'global':
                query = query.filter(AppleMusicContainerStreams.country_code == country_code)

            session.add_all([
                AppleWeeklyTopPlaylist(
                    playlist_id=q.container_id,
                    country_code=country_code,
                    date=week_start_date,
                    rank=idx
                ) for idx, q in enumerate(query[:TOP_PLAY_LISTS_NUMBER], start=1)
            ])
            session.commit()
            logger.info(f'Date: {week_start_date} - {country_code} - {time.time() - start_time}')

    logger.info(f'Done')


if __name__ == '__main__':
    start_date = datetime.strptime(os.environ['START_DATE'], '%Y-%m-%d').date()
    end_date = datetime.strptime(os.environ['END_DATE'], '%Y-%m-%d').date()
    back_fill_data(start_date, end_date)