from datetime import timedelta
from typing import List

import athena
import config
import constants.common as consts
import main_db
import s3
import utils
from logger import logger
from metric_data import MetricData, get_metric_data


def load_dates(metric_data: MetricData):
    """Load dates to S3.

    Args:
        metric_data: Metric data.

    """
    data = [
        (config.DATE_FROM + timedelta(days=i),) for i in range((metric_data.latest_date - config.DATE_FROM).days + 1)
    ]
    s3_path = s3.get_relative_path(consts.Table.DATE, consts.Folder.UPLOAD)
    s3.write_csv(s3_path + "dates.csv", data)


def load_playlists(playlist_ids: List[str]):
    """Load playlist ID list to S3.

    Args:
        playlist_ids: Playlist ID list.

    """
    s3_path = s3.get_relative_path(consts.Table.PLAYLIST, consts.Folder.UPLOAD)
    s3.write_csv(s3_path + "playlists.csv", [(i,) for i in playlist_ids])


def create_tables(table_type: str):
    """Create all tables.

    Args:
        table_type: Table type.

    """
    if config.ATHENA_TRY_CREATE_DATABASE:
        athena.create_database()

    athena.create_table(table_type)
    athena.create_table(consts.Table.DATE)
    athena.create_table(consts.Table.PLAYLIST)


def drop_tables(table_type: str):
    """Drop tables.

    Args:
        table_type: Table type.

    """
    athena.drop_table(table_type)
    athena.drop_table(consts.Table.DATE)
    athena.drop_table(consts.Table.PLAYLIST)


def clean_tables(table_type: str, with_playlists: bool = True):
    """Clear tables.

    Args:
        table_type: Table type.
        with_playlists: Clean playlists as well.

    """
    s3.delete_folder_files(s3.get_relative_path(table_type, consts.Folder.UPLOAD))
    if with_playlists:
        s3.delete_folder_files(s3.get_relative_path(consts.Table.PLAYLIST, consts.Folder.UPLOAD))
    s3.delete_s3_files_except_csv(s3.get_relative_path(table_type, consts.Folder.DOWNLOAD))


@utils.timing
def process_streams(
    table_type: str, metric_data: MetricData, playlist_chunk: List[str], missing_playlists: List[str], index: int
):
    """Process streams data.

    Args:
        table_type: Table type.
        metric_data: Metric data.
        playlist_chunk: Playlist ID list chunk.
        missing_playlists: Missing playlists.
        index: Chunk index.

    """
    result = 0

    log_playlist = f"{index}/{len(missing_playlists)}"

    for i, country_code in enumerate(consts.COUNTRY_CODE_LIST):
        logger.debug(f"{table_type}: {log_playlist} {country_code} ({i + 1}/{len(consts.COUNTRY_CODE_LIST)}) uploading")

        for j, distributer_id in enumerate(consts.DISTRIBUTER_ID_LIST):
            logger.debug(f"{table_type}: {distributer_id} ({j + 1}/{len(consts.DISTRIBUTER_ID_LIST)})")
            main_db.upload_to_s3(
                table_type,
                s3.get_full_path(table_type, consts.Folder.UPLOAD),
                f"{country_code}_{distributer_id}",
                playlist_id_list=tuple(playlist_chunk),
                date_from=config.DATE_FROM - timedelta(days=13),
                date_to=metric_data.latest_date,
                distributer_id=distributer_id,
                country_code=country_code,
            )

    logger.debug(f"{table_type}: {log_playlist} aggregating")
    athena.execute_aggregate_query(table_type)

    clean_tables(table_type, with_playlists=True)

    logger.debug(f"{table_type}: {log_playlist} downloading")
    result += main_db.download_from_s3(table_type, s3.get_full_path(table_type, consts.Folder.DOWNLOAD))

    s3.delete_folder_files(s3.get_relative_path(table_type, consts.Folder.DOWNLOAD))

    return result


def process_followers(
    table_type: str, metric_data: MetricData, playlist_chunk: List[str], missing_playlists: List[str], index: int
):
    """Process followers data.

    Args:
        table_type: Table type.
        metric_data: Metric data.
        playlist_chunk: Playlist ID list chunk.
        missing_playlists: Missing playlists.
        index: Chunk index.

    """
    logger.debug(f"{table_type}: {index}/{len(missing_playlists)} uploading")
    main_db.upload_to_s3(
        table_type,
        s3.get_full_path(table_type, consts.Folder.UPLOAD),
        playlist_id_list=tuple(playlist_chunk),
        date_from=config.DATE_FROM - timedelta(days=14),
        date_to=metric_data.latest_date,
    )

    logger.debug(f"{table_type}: {index}/{len(missing_playlists)} aggregating")
    athena.execute_aggregate_query(table_type)

    clean_tables(table_type)

    logger.debug(f"{table_type}: {index}/{len(missing_playlists)} downloading")
    result = main_db.download_from_s3(table_type, s3.get_full_path(table_type, consts.Folder.DOWNLOAD))

    s3.delete_folder_files(s3.get_relative_path(table_type, consts.Folder.DOWNLOAD))

    return result


@utils.timing
def add_missing_data(table_type: str, metric_data: MetricData, missing_playlists: List[str]) -> int:
    """Add missing playlists history data.

    Args:
        table_type: Update table type.
        metric_data: Previous and latest dates.
        missing_playlists: Missing playlists.

    Returns:
        int: Rows count.

    """
    if not missing_playlists:
        return 0

    result = 0

    process_data = process_streams if table_type == consts.Table.STREAMS else process_followers

    chunk_size = config.PLAYLIST_AGGREGATE_METRIC_CHUNK_SIZE[table_type]

    try:
        create_tables(table_type)

        load_dates(metric_data)

        for index in range(0, len(missing_playlists), chunk_size):
            playlist_chunk = missing_playlists[index : index + chunk_size]

            load_playlists(playlist_chunk)

            process_data(table_type, metric_data, playlist_chunk, missing_playlists, index)

            count = metric_data.db.insert_metric_playlists(playlist_chunk, metric_data.latest_date)
            logger.debug(f"{table_type}: {index}/{len(missing_playlists)} {count} inserted")

    finally:
        clean_tables(table_type)
        s3.delete_folder_files(s3.get_relative_path(table_type, consts.Folder.DOWNLOAD))
        s3.delete_folder_files(s3.get_relative_path(consts.Table.DATE, consts.Folder.UPLOAD))
        drop_tables(table_type)
        s3.delete_temp()

    return result


def aggregate_data() -> dict:
    """Aggregate missing playlists data.

    Returns:
        Table type to full new or not.

    """
    full_new = {consts.Table.STREAMS: False, consts.Table.FOLLOWERS: False}
    if config.AGGREGATE_MISSING_PLAYLISTS:
        for table_type in config.UPDATE_TABLE_LIST:
            metric_data = get_metric_data(table_type)

            top_playlists = metric_data.db.get_top_playlists()
            logger.debug(f"{table_type}: {len(top_playlists)} top")
            aggregated_playlists = metric_data.db.get_aggregated_playlists()
            logger.debug(f"{table_type}: {len(aggregated_playlists)} aggregated")
            if config.UPDATE_BY_AGGREGATING:
                missing_playlists = aggregated_playlists
            else:
                missing_playlists = list(set(top_playlists) - set(aggregated_playlists))

            if missing_playlists:
                logger.debug(f"{table_type}: {len(missing_playlists)} missing")
                row_count = add_missing_data(table_type, metric_data, missing_playlists)
                logger.debug(f"{table_type}: {row_count} new records added")

            full_new[table_type] = not aggregated_playlists

    return full_new