import datetime
import logging

import pandas as pd

from tadas.domain import constants
from tadas.snowflake import client as snowflake_utils
from tadas.platform import config

logger = logging.getLogger(__name__)

SUPPORTED_DSPS = ('spotify', 'apple', 'amazon', 'tiktok', 'youtube')

DSP_FILTERS = {
    "spotify": "and store_id=286 and feed_id=1",
    "apple": "and store_id=1 and feed_id=4",
    "amazon": "and store_id=187 and feed_id in (35, 36, 37)",
    "youtube": "and store_id=453 and feed_id=38",
    "tiktok": "",
}
DSP_SOURCE = {
    "spotify": "v_streams_by_track_country_feed_distributor_daily",
    "apple": "v_streams_by_track_country_feed_distributor_daily",
    "amazon": "v_streams_by_track_country_feed_distributor_daily",
    "youtube": "v_streams_by_track_country_feed_distributor_daily",
    "tiktok": "v_tiktok_by_isrc_country_daily",
}


def is_table_fresh(table_name, expected_date: datetime.date, date_column='report_date'):
    """
    Return True if table exists and has some data for given date. False otherwise.

    :param table_name:
    :param expected_date:
    :param date_column:
    :return:
    """

    with snowflake_utils.snowflake_connection() as conn:
        if not snowflake_utils.is_table_exists(table_name, conn):
            return False

        query = f'''
        select count(*) 
        from {table_name}
        where {date_column} = %(expected_date)s
        '''
        params = {'expected_date': expected_date}
        results = snowflake_utils.query_snowflake_fetchall(
            connection=conn,
            query=query,
            params=params,
        )
        if not results:
            return False

        found_rows = results[0][0]
        logger.info(f'Found {found_rows} in {table_name} for {expected_date}')
        return found_rows > 0


def load_available_dates(
        max_days_back=5,
        distributors=('sme', 'theorchard'),
        dsps=SUPPORTED_DSPS,
):
    """Check which days have data looking up to given days back."""
    logger.info(f'Checking data availability for {max_days_back=}')

    source_schema = config.get('SOURCE_SCHEMA')
    available_dates = {}
    with snowflake_utils.snowflake_connection() as conn:
        for distributor in distributors:
            for dsp in dsps:
                if dsp == 'tiktok':
                    # tiktok DBT models don't keep licensor/distributor info.
                    # but as soon as DBT for tiktok is build only when both sme and theorchard are available,
                    # we can assume that data is available for both.
                    distributor_filter = ''
                else:
                    distributor_filter = f'and distributor = \'{distributor}\''

                dsp_filter = DSP_FILTERS[dsp]
                dsp_source = DSP_SOURCE[dsp]
                query = f"""select download_activity_date as report_date, count(*) as sum_streams
                    from {source_schema}{dsp_source}
                    where download_activity_date >= current_date() - {max_days_back}
                        {distributor_filter}
                        {dsp_filter}
                    group by download_activity_date
                    order by download_activity_date desc
                """
                logger.info(f'Checking data availability for {dsp} ({distributor}) on {max_days_back} days back.')
                params = {'max_days_back': max_days_back}
                results_key = f'{dsp}_{distributor}'
                data_for_results_key = snowflake_utils.query_snowflake_fetchall(
                    connection=conn,
                    query=query,
                    params=params
                )
                logger.info(f'Available data for {results_key} = {data_for_results_key}')
                # filter only dates with streams
                available_dates[results_key] = [
                    r[0]
                    for r in data_for_results_key
                    if r[1] > 0
                ]

    for key, dates in available_dates.items():
        logger.info(f'Available dates for {key}: {dates}')

    return available_dates


def prepare_sql_select_from_dbt_model(
        source_table: str,
        tracks_table: str,
        columns_definition: dict,
        periods: dict,
        group_by_country=True) -> str:
    columns_expressions = []
    dates_expressions = []
    source_schema = config.get('SOURCE_SCHEMA')

    for period, period_props in periods.items():
        dates_expressions.append(f'{period_props["date_range"]}')
        for column, col_props in columns_definition.items():
            if 'condition' in col_props:
                condition = f' and {col_props["condition"]}'
            else:
                condition = ''

            prefix = period_props['column_prefix']
            col_select_expression = (f'sum('
                                     f'case '
                                     f'when download_activity_date = {period_props["date_range"]}{condition} '
                                     f'then {col_props["column"]} '
                                     f'else 0 '
                                     f'end) '
                                     f'as {prefix}{column}')
            columns_expressions.append(col_select_expression)

    sp_q_t = f"""
    select 
        t.isrc_cd,
        t.geo_country as geo_country,
        {',\n'.join(columns_expressions)}
    from {source_schema}{source_table} f
    join {tracks_table} t on f.isrc = t.isrc_cd {'and f.country_code = t.geo_country' if group_by_country else '' }
    where 
     download_activity_date in ({', '.join(dates_expressions)})
    group by all
            """

    return sp_q_t


def select_tracks(report_date, group_by_country: bool):
    # from backup/tadas_data_generation*.py
    # file           | agg_streams_week      | agg_streams_yest     |
    # ---------------|-----------------------|----------------------|
    # orchard        | 35000 (spotify,apple) | 5000 (spotify,apple) |
    # orchard (big geos) | 35000 (spotify,apple) | 5000 (spotify,apple) |
    # isrc           | 70000 (spotify+apple) | 5000 (spotify+apple) |
    # isrc (big geos)| 140000 (spotify+apple)| 10000 (spotify+apple)|
    #### Globals
    # orchard global | 35000 (spotify,apple) | 5000 (spotify,apple) |
    # isrc global    | 70000 (spotify+apple) | 5000 (spotify+apple) |
    if group_by_country:
        agg_streams_week = 70_000
        agg_streams_yest = 10_000
        geo_country_expr = 'f.country_code'
    else:
        agg_streams_week = 70_000
        agg_streams_yest = 10_000
        geo_country_expr = "'XX'"

    geos = constants.geos
    source_schema = config.get('SOURCE_SCHEMA')

    params = {
        'report_date': report_date,
        'agg_streams_yest': agg_streams_yest,
        'agg_streams_week': agg_streams_week,
        'countries': geos,
    }
    # Pull tracks with agg'ed above limit over past week
    query_agg_streams_week = f""" 
        SELECT f.isrc as isrc_cd, 
                {geo_country_expr} as geo_country,
                sum(f.streams) as agg_streams
        FROM {source_schema}v_streams_by_track_country_feed_distributor_daily f
        WHERE f.download_activity_date between %(report_date)s::date - 7 and %(report_date)s::date - 1
                and f.country_code in (%(countries)s)
        GROUP BY 1,2
        HAVING agg_streams >= %(agg_streams_week)s
        """
    # pull tracks with above limit for yesterday.
    query_agg_streams_yest = f""" 
        SELECT f.isrc as isrc_cd, 
                {geo_country_expr} as geo_country,
                sum(f.streams) as yest_streams
        FROM {source_schema}v_streams_by_track_country_feed_distributor_daily f
        WHERE f.download_activity_date = %(report_date)s::date
                and f.country_code in (%(countries)s)
        GROUP BY 1,2
        HAVING yest_streams >= %(agg_streams_yest)s
    """

    with snowflake_utils.snowflake_connection() as conn:
        sp_legacy_by_geo_2w_df = snowflake_utils.query_snowflake_to_df(
            query=query_agg_streams_week,
            params=params,
            connection=conn,
        )

        legacy_by_geo_2w_df = sp_legacy_by_geo_2w_df
        legacy_by_geo_2w_df = legacy_by_geo_2w_df.groupby(['isrc_cd', 'geo_country']).sum().reset_index()

        sp_legacy_by_geo_yest_df = snowflake_utils.query_snowflake_to_df(
            query=query_agg_streams_yest,
            params=params,
            connection=conn,
        )

    legacy_by_geo_yest_df = sp_legacy_by_geo_yest_df
    legacy_by_geo_yest_df = legacy_by_geo_yest_df.groupby(['isrc_cd', 'geo_country']).sum().reset_index()

    # merge the versions together to get the outer list.
    legacy_by_geo = legacy_by_geo_2w_df.merge(legacy_by_geo_yest_df, how='outer', on=['isrc_cd', 'geo_country'])
    orchard_by_geo = legacy_by_geo

    # filter up for geos not
    # TODO: use common.big_geos after refactoring
    big_geos = ['US', 'FR', 'MX', 'BR', 'DE', 'GB', 'IN', 'PH', 'CN']
    big_geos_df = orchard_by_geo[orchard_by_geo['geo_country'].isin(big_geos)]
    small_geos_df = orchard_by_geo[~orchard_by_geo['geo_country'].isin(big_geos)]

    big_geos_df['is_big_avg'] = 0
    big_geos_df['is_big_avg'][big_geos_df['agg_streams'] > agg_streams_week] = 1
    big_geos_df['is_big_yest'] = 0
    big_geos_df['is_big_yest'][big_geos_df['yest_streams'] > agg_streams_yest] = 1
    big_geos_df['is_combo'] = big_geos_df['is_big_avg'] + big_geos_df['is_big_yest']
    big_geos_df = big_geos_df[big_geos_df['is_combo'] > 0]

    orchard_by_geo = pd.concat([big_geos_df, small_geos_df], ignore_index=True)

    # make final version with only columns needed
    tracks_for_legacy = orchard_by_geo[['isrc_cd', 'geo_country']]

    # I'm going to temporarily remove Nashville and RCA data from TADAS
    tracks_df = tracks_for_legacy.copy()
    tracks_df['pfn_geo'] = tracks_df['isrc_cd'].astype(str) + '_' + tracks_df['geo_country'].astype(str)

    tracks_df['report_date'] = pd.to_datetime(report_date, format="%Y-%m-%d")

    return tracks_df