#!/usr/bin/python3

from lib.covidData import covidData


def set_data_dict(selection, import_data):
    columns = [i for i in covidData.__dict__.keys() if i[:1] != '_'][1:]
    formatted_import_data = [d.strip().replace(
        ',', '').replace('+', '') for d in import_data]
    output_dict = {
        'selection': selection.upper().strip().replace('\n', ''),
        'selection_original': selection.strip().replace('\n', '')
    }
    for i in range(1, len(columns)):
        if(formatted_import_data[i] and not formatted_import_data[i] == ''):
            if(i >= len(columns) - 2):
                output_dict[columns[i]] = formatted_import_data[i]
            else:
                try:
                    output_dict[columns[i]] = int(formatted_import_data[i])
                except:
                    output_dict[columns[i]] = -1
        else:
            output_dict[columns[i]] = -1
    return output_dict


def get_data_dict(query_data):
    columns = [i for i in covidData.__dict__.keys() if i[:1] != '_']
    output_dict = {
        'selection': query_data.selection,
        'selection_original': query_data.selection_original
    }
    for i in range(2, len(columns)):
        if(i >= len(columns) - 2):
            output_dict[columns[i]] = query_data.__getattribute__(
                columns[i])
        else:
            output_dict[columns[i]] = int(
                query_data.__getattribute__(columns[i]))
    return output_dict


def set_data(session, selection, import_data):
    new_data = covidData()
    formatted_data_dict = set_data_dict(selection, import_data)
    try:
        for key in formatted_data_dict:
            new_data.__setattr__(key, formatted_data_dict[key])
    except Exception as e:
        print(e)
    session.merge(new_data)
    session.commit()


def get_data(session, selection):
    print('Formatting data.')
    all_data_query = session.query(covidData).filter(
        covidData.selection == selection).all()
    return get_data_dict(all_data_query[0])


def get_top_n_rows(session, num):
    print(f'Getting top {num} rows.')
    filter_selections = ['WORLD', 'TOTAL', 'EUROPE', 'NORTH AMERICA', 'ASIA', 'SOUTH AMERICA', 'AFRICA', 'OCEANIA']
    exception_selection = session.query(covidData).filter(
        covidData.selection.in_(filter_selections))
    exception_null = session.query(covidData).filter(covidData.total_cases == -1)
    top_n_rows = session.query(covidData).except_all(exception_null, exception_selection).order_by(
        covidData.total_cases.desc()).limit(num).all()
    return [get_data_dict(n) for n in top_n_rows]