import csv
import json
import itertools
import os

from dotenv import load_dotenv

from snowflake_connector.etl_connector import SnowflakeSQLExecutor
from snowflake_connector.metadata_connector import SnowflakeMetadataConnector, SnowflakeBadResponse

load_dotenv()

SF_PARAMS = {
    'role': os.getenv('ROLE'),
    'warehouse': os.getenv('WAREHOUSE'),
    'db': os.getenv('DB'),
    'schema': os.getenv('SCHEMA'),
    'user': os.getenv('USER'),
    'password': os.getenv('PASSWORD'),
    'account': os.getenv('ACCOUNT')}

SCENARIOS = [
    {
        'params': [
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205','BRRGE1508687','BRRGE1707681'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205','BRRGE1508687'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702'"},
            {'isrcs': "'BRRGE1708139'"},
        ],
        'sql_template':
            '''
            SELECT
                random(),
                territory_code,
                --hll_estimate(hll_combine(hll_import(hll_users_variant))) AS listeners_number,
                hll_estimate(hll_combine(hll_import(parse_json(hll_users_varchar)))) as listeners_number,
                SUM(streams) AS streams_number
            FROM DEV_ENGINEERING.NSNYDER.summary_geographics
            WHERE download_activity_date BETWEEN '2017-02-28' AND '2018-02-28' and
              labelid = 21786
              and isrc in ({isrcs})
              --and artistid in (898003)--,571615,503140,544572,740512,557298)
            GROUP BY territory_code
            ORDER BY streams_number DESC;
            '''
    },
    {
        'params': [
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205','BRRGE1508687','BRRGE1707681'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205','BRRGE1508687'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001','BRRGE1708205'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702','BXBAL1700001'"},
            {'isrcs': "'BRRGE1708139','BRRGE1508702'"},
            {'isrcs': "'BRRGE1708139'"},
        ],
        'sql_template':
            '''
            SELECT
                random(),
                territory_code,
                hll_estimate(hll_combine(hll_import(hll_users_variant))) AS listeners_number,
                --hll_estimate(hll_combine(hll_import(parse_json(hll_users_varchar)))) as listeners_number,
                SUM(streams) AS streams_number
            FROM DEV_ENGINEERING.NSNYDER.summary_geographics
            WHERE download_activity_date BETWEEN '2017-02-28' AND '2018-02-28' and
              labelid = 21786
              and isrc in ({isrcs})
              --and artistid in (898003)--,571615,503140,544572,740512,557298)
            GROUP BY territory_code
            ORDER BY streams_number DESC;
            '''
    }]

FILENAME = 'results.csv'

REQUESTED_FIELDS = [
    'status',
    'warehouseName',
    'startTime',
    'endTime',
    'totalDuration',
    'stats.scanBytes',
    'stats.producedRows',
    'id',
    'queryTag',
    'sqlTextHash',
    'sqlText',
    'stats.ioLocalFdnReadBytes',
    'stats.ioRemoteFdnReadBytes'
]


def get_key(data, key):

    def helper(data, path):
        data_part = data.get(path[0])
        if len(path) == 1:
            return data_part
        else:
            return helper(data_part, path[1:])

    return helper(data, key.split('.'))


def flatten_query(query_data):
    return dict((key, get_key(query_data, key)) for key in REQUESTED_FIELDS)


def format_stats(data):
    return list(map(flatten_query, data.get('queries')))


def write(data, format_func):
    with open(FILENAME, 'w') as output:
        format_func(output, data)


def as_csv(fd, data):
    keys = data[0].keys()
    dict_writer = csv.DictWriter(fd, keys)
    dict_writer.writeheader()
    dict_writer.writerows(data)


def as_json(fd, data):
    json.dump(data, fd)


def get_query_data(sfqid):
    meta = SnowflakeMetadataConnector(SF_PARAMS)
    meta.authenticate()
    return meta._get_query_data(sfqid)


def run_query_and_report(query):
    with SnowflakeSQLExecutor(SF_PARAMS) as executor:
        cursor = executor.execute(query)

    return get_query_data(cursor.sfqid)


def bind_params(scenario):
    sql_template = scenario['sql_template']
    params = scenario['params']
    return [sql_template.format(**param_set) for param_set in params]


def flatten(nested_lists):
    return list(itertools.chain(*nested_lists))


if __name__ == "__main__":

    '''
    # Use this pattern to run scenarios and get results.
    queries = flatten(bind_params(scenario) for scenario in SCENARIOS)
    stats = [run_query_and_report(query) for query in queries]
    '''
    # Use this pattern to provide existing query IDs and get statss.
    stats = [get_query_data(sfqid) for sfqid in ['49867132-a931-4145-b160-58c9585244e8']]
    formatted_stats = flatten(format_stats(stat_set) for stat_set in stats)
    write(formatted_stats, as_csv)