import os
import random
from timeit import default_timer as timer
import numpy as np
from snowflake_connector.etl_connector import SnowflakeSQLExecutor
import json

label_ids = [
    9513,
    11233,
    11382,
    11454,
    12239,
    12602,
    12643,
    12650,
    13928,
    13977,
    13671,
    28676,
    28668,
    29195,
    29251,
    29339,
    22266,
    22334,
    22340,
    23517,
    23540,
    23581,
    23757,
    23806,
    23930,
    24090,
    24110,
    25519,
    20156,
    20178,
    20266,
    19185,
    17422,
    17477,
    18000,
    22553,
    15320,
    16971,
    17069,
    25670,
    25690,
    19231,
    19239,
    19383,
    22399,
    11667,
    12051,
    15494,
    21795,
    8631,
    8851,
    24329,
    196,
    1258,
    1726,
    1796,
    1463,
    2179,
    2638,
    2668,
    3386,
    4759,
    3611,
    4701,
    5249,
    5404,
    6012,
    6119,
    18116,
    18667,
    14119,
    14704,
    14708,
    14741,
    17204,
    17307,
    21945,
    9081,
    9134,
    9355,
    21693,
    21740,
    18991,
    18992,
    19079,
    19093,
    19103,
    18150,
    18179,
    18212,
    22778,
    22844,
    23374,
    26994,
    27015,
    27162,
    27185,
    27203,
    27453,
    27456,
]

isrcs = [
    'USCRS0691277',
    'USCRS0692116',
    'USCRS1092022',
    'USCRS0896503',
    'USCRS0997120',
    'USCRS1091426',
    'USCRS1295801',
    'USCRS1493812',
    'QMFME1349238',
    'USRS72600001',
    'US4CC0410043',
    'USDAZ0510702',
    'USDAZ0511849',
    'US4A70400292',
    'US4J51100033',
    'USAK10110855',
    'QM6MZ1772235',
    'ITN581300122',
    'GBHJV0810018',
    'AUTQ00100027',
    'AUTQ00600075',
    'AUTQ00600079',
    'AUTQ00800211',
    'AUTQ01000033',
    'US4TX1200064',
    'US4QA1100013',
    'USAHQ0510003',
    'USAHQ0910118',
    'QMFME1455303',
    'QMFMG1466303',
    'USHTV0810098',
    'USHTV0910104',
    'USHTV1100085',
    'ITN581300035',
    'US4CM0910202',
    'USHTV0810039',
    'US78L1010004',
    'GBCHV0150471',
    'DK6KA1300913',
    'JOA021300115',
    'US2V91010034',
    'US54H0312310',
    'US54H0312318',
    'QMBZ91585062',
    'USDMG0312702',
    'USDMG0313713',
    'USDMG0514709',
    'US4BQ0401270',
    'US4BQ0401639',
    'US76A0510089',
    'US75J0910109',
    'CAI371010506',
    'CAI371114213',
    'US25T0616635',
    'USBBW0802806',
    'GBAJC9902191',
    'SEVCK0604003',
    'US4BP0820202',
    'USA310727209',
    'USA310726405',
    'USRHR8400203',
    'GBBDZ1026461',
    'US4A70400101',
    'US4A70400132',
    'AUTQ00300041',
    'AUTQ00700219',
    'AUTQ01000031',
    'US4TX1200082',
    'US4QA0500150',
    'QMFME1455300',
    'USHTV1010168',
    'USHTV1010217',
    'QM6P41440230',
    'USHTV1100688',
    'US52U0410091',
    'USZUR0602098',
    'QMVTG1200105',
    'US5240610013',
    'US78L0510138',
    'US35K0490014',
    'USL4R0782244',
    'DK6KA1400102',
    'DKEY20901101',
    'DKPC91100303',
    'USJ7U0842166',
    'US54H0312320',
    'US54H0513280',
    'USDMG0312717',
    'USDMG0313110',
    'US3Q50710035',
    'US4BQ0400404',
    'US4BQ0401637',
    'CA1M61000001',
    'CAI370713416',
    'CAI370913320',
    'CAI371679223',
    'US25T0523749',
    'US25T0523960',
    'US6531010001',
    'USBBW1000067',
]

sf_config = {
    'role': os.environ.get('SNOWFLAKE_ROLE'),
    'warehouse': os.environ.get('SNOWFLAKE_WAREHOUSE'),
    'db': os.environ.get('SNOWFLAKE_DATABASE'),
    'schema': os.environ.get('SNOWFLAKE_SCHEMA'),
    'user': os.environ.get('SNOWFLAKE_USER'),
    'password': os.environ.get('SNOWFLAKE_PASSWORD'),
    'account': os.environ.get('SNOWFLAKE_ACCOUNT')
}


def query(label_only=True, limit=10):
    if label_only:
        clause = 'TRUE'
    else:
        clause = 'isrc = %(isrc)s'
    with SnowflakeSQLExecutor(sf_config) as executor:
        with executor.get_cursor() as cursor:
            session_id = executor.snowflake_conn.session_id
            count = 0
            sql = """
                SELECT * FROM placements_test
                WHERE labelid = %(labelid)s AND {clause} 
                LIMIT %(limit)s;
            """.format(clause=clause)
            cursor.execute(sql, {
                'labelid': random.choice(label_ids),
                'isrc': random.choice(isrcs),
                'limit': limit,
            })
            for _ in cursor:
                count += 1
            return session_id, cursor.sfqid, count


def get_sys_time(session_id, query_id):
    with SnowflakeSQLExecutor(sf_config) as executor:
        sys_time, *_ = executor.fetchone("""
            SELECT total_elapsed_time
            FROM TABLE(information_schema.query_history_by_session(
                session_id=>%(session_id)s)
            )
            WHERE query_id = %(query_id)s AND
                  execution_status = 'SUCCESS';
        """, {'session_id': session_id, 'query_id': query_id})
        return float(sys_time) / 1000.0


def calc_metrics(l):
    return {
        'avg': np.average(l),
        'median': np.median(l),
        '95p': np.percentile(l, 95),
        'max': max(l)
    }


def analyze(label_only=True, n=13):
    print('Label only: {}'.format(label_only))
    print('WH: {}'.format(sf_config['warehouse']))

    usr_times = []
    sys_times = []
    for i in range(n):
        start = timer()
        session_id, query_id, row_cnt = query(label_only)
        usr_time = timer() - start
        sys_time = get_sys_time(session_id, query_id)
        print('{}: (usr: {:.3f}; sys: {:.3f}) [Rows: {}]'.format(
            session_id, usr_time, sys_time,  row_cnt))
        usr_times.append(usr_time)
        sys_times.append(sys_time)

    print('Usr time: {}'.format(json.dumps(calc_metrics(usr_times), indent=4)))
    print('Sys time: {}'.format(json.dumps(calc_metrics(sys_times), indent=4)))


def main():
    analyze(label_only=True)
    analyze(label_only=False)


if __name__ == '__main__':
    main()