from airflow.models import Variable
import requests
import json
from airflow.contrib.operators.ssh_operator import SSHOperator

# This is to get the spotify access token for a licensor
def get_access_token(licensor):
    oauth_url = Variable.get('spotify_oauth_base_api_url') + '&client_id=' + \
                Variable.get(licensor['client_id_param']) + '&client_secret=' + \
                Variable.get(licensor['client_secret_param'])
    print('Before calling OAuth url')
    r = requests.post(url = oauth_url)
    print('status ==', r.status_code)
    access_token = r.json()['access_token']
    return access_token

# This is to get licensor details for a given licensor name
def get_licensor_details(name):
    licensor = None
    licensors = json.loads(Variable.get('licensors'))
    for ref_licensor in licensors:
        if ref_licensor['name'] == name:
            licensor = ref_licensor
            break
    return licensor

# This is to check if the downloads are available for a given date
def are_downloads_available(licensor_name, date_to_download):
    licensor = get_licensor_details(licensor_name)
    access_token = get_access_token(licensor)
    available_countries_streams_url = Variable.get('spotify_base_api_url') + '/' + \
                                  licensor['name'] + 'v2/streams/' + date_to_download.replace('-', '/') + \
                                  '?oauth_token=' +  access_token

    r = requests.get(available_countries_streams_url)
    available_country_streams = r.json();
    downloads_available = False
    if len(available_country_streams) > 0:
        downloads_available = True

    return downloads_available

# This function is for creating the download tasks and it spreads out the SSH tasks based on the country
def create_download_tasks(date_to_download, dag, fork_dwnlds, wait_task, end_task):
    licensors = json.loads(Variable.get('licensors'))
    # Download the extracted files and spread out the downloads across 4 VMs
    for licensor in licensors:
        if licensor['name'] == 'sonybmgmusicentertainment':
            # 1 - (Only SonyBMG) This task creates download for only US stream file
            countries_include = '["US"]'
            licensor_task1 = SSHOperator(
                ssh_conn_id='remote_vm_conn',
                task_id=licensor['name'] + '_' + '_dwnlds_US',
                command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \
                        ' -dd ' + date_to_download + ' -dt extracted_files -ci \'' + countries_include + \
                        '\' -dncf False',
                dag=dag)

            licensor_task1.set_upstream(fork_dwnlds)
            licensor_task1.set_downstream(wait_task)

            # 2 - (Only SonyBMG) This task creates download for countries that have large stream files
            countries_include = '["GB", "SE", "DE",  "BR", "TR", "PL", "NL", "MX"]'
            licensor_task2 = SSHOperator(
                ssh_conn_id='remote_vm_conn1',
                task_id=licensor['name'] + '_' + '_dwnlds_GB_SE_DE_BR_TR_PL_NL_MX',
                command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \
                        ' -dd ' + date_to_download + ' -dt extracted_files -ci \'' + countries_include + \
                        '\' -dncf False',
                dag=dag)
            licensor_task2.set_upstream(fork_dwnlds)
            licensor_task2.set_downstream(wait_task)

            # 3 - (Only SonyBMG) This task creates download for rest of the countries
            #     and all files including tracks, users & agg streams
            countries_exclude = '["US", "MX", "NL", "GB", "SE", "DE",  "BR", "TR", "PL"]'
            licensor_task3 = SSHOperator(
                ssh_conn_id='remote_vm_conn2',
                task_id=licensor['name'] + '_' + '_dwnlds_REST_COUNTRIES',
                command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \
                        ' -dd ' + date_to_download + ' -dt extracted_files -ce \'' + countries_exclude + \
                        '\' -dncf True',
                dag=dag)
            licensor_task3.set_upstream(fork_dwnlds)
            licensor_task3.set_downstream(wait_task)
        else:
            # 4 - (Orchard, sonymarketing & intl) This task creates download for all countries
            #     and all files including tracks, users & agg streams
            licensor_task = SSHOperator(
                ssh_conn_id='remote_vm_conn3',
                task_id=licensor['name'] + '_' + '_dwnlds_ALL',
                command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \
                        ' -dd ' + date_to_download + ' -dt extracted_files',
                dag=dag)
            licensor_task.set_upstream(fork_dwnlds)
            licensor_task.set_downstream(wait_task)

    # Download the archive files and spread out the downloads across 4 VMs
    lc_idx = 0
    for licensor in licensors:
        connection_id = 'remote_vm_conn'
        if lc_idx > 0:
            connection_id = 'remote_vm_conn' + str(lc_idx)

        licensor_task = SSHOperator(
            ssh_conn_id=connection_id,
            task_id=licensor['name'] + '_' + '_gz_dwnlds',
            command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \
                    ' -dd ' + date_to_download + ' -dt gzip_files',
            dag=dag)
        licensor_task.set_upstream(wait_task)
        licensor_task.set_downstream(end_task)