"""Util functions for download_to_db garcon task."""

import datetime

from flows.theatrical import utils


GROSS_COLUMNS = [
    'mon_gross', 'tue_gross', 'wed_gross', 'thu_gross', 'fri_gross',
    'sat_gross', 'sun_gross']

SOURCE_COLUMNS = [
    'company_name', 'film_id', 'upc', 'film_name', 'division_id',
    'division_name', 'branch_id', 'branch_name', 'circuit_id', 'theater_id',
    'theater_name', 'city_state', 'booking_type', 'booking_no', 'week_no',
    'play_date', 'gross_type', 'fri_pcode', 'sat_pcode', 'sun_pcode',
    'mon_pcode', 'tue_pcode', 'wed_pcode', 'thu_pcode', 'is_fri_est',
    'is_sat_est', 'is_sun_est', 'is_mon_est', 'is_tue_est', 'is_wed_est',
    'is_thu_est', 'curr_3day', 'curr_7day', 'prev_3day', 'prev_7day',
    'mkt_name', 'dma_id', 'dma_name', 'ranking']
NEW_COLUMNS = ['date', 'gross']
CSV_STRUCTURE = SOURCE_COLUMNS + NEW_COLUMNS


def gross_columns_for_period(date_start, date_end):
    """Select needed gross columns and calculate corresponding dates.

    Args:
        date_start (date): start period for the file.
        date_end (date): end period for the file.

    Returns:
        dict: Dict of {gross_column: corresponding_date}.
    """
    # pruning incorrect periods or periods longer than a week
    if date_start > date_end or (date_end - date_start).days + 1 > 7:
        return {}

    weekday_start = date_start.weekday()
    weekday_end = date_end.weekday()

    # in case of overflowing
    if weekday_end <= weekday_start and date_start < date_end:
        weekday_end += 7

    res = {}
    for weekday in range(weekday_start, weekday_end + 1):
        res.update(
            {GROSS_COLUMNS[weekday % 7]: date_start + datetime.timedelta(
                days=weekday - weekday_start)})
    return res


def filter_and_transform(header, rows, date_start, date_end):
    """Transform theatrical raw data.

    Args:
        header (list(str)): list of column headers.
        rows (list(list(str))): list of CSV file rows.
        date_start (date): first day needed from given file.
        date_end (date): last day needed from given file.

    Returns:
        tuple: new header, new rows, and upcs in the rows.
    """
    gross_columns = gross_columns_for_period(date_start, date_end)
    gross_indexes = {key: header.index(key) for key in gross_columns.keys()}

    upc_index = header.index('upc')

    new_rows = []
    new_upcs = set()

    old_indexes = {h: header.index(h) for h in SOURCE_COLUMNS}

    for row in rows:
        new_upcs.add(row[upc_index])
        for gross_column, date in gross_columns.items():
            new_row = []
            for h in SOURCE_COLUMNS:
                col = row[old_indexes[h]]
                new_row.append(col.strip())
            row_date = utils.serialize_date(date)
            row_gross = row[gross_indexes[gross_column]]
            row_gross = row_gross.strip()
            row_gross = row_gross.replace(',', '')
            new_row.append(row_date)
            new_row.append(row_gross)
            new_rows.append(new_row)
    return CSV_STRUCTURE, new_rows, tuple(new_upcs)