import os
import random
from datetime import date, timedelta
from locust import HttpLocust, TaskSet, task
from http import client
import json

urls = {
    'artist': '/screens/artist/:sonyArtistId?startDate=:startDate&endDate=:endDate',
    'countries': '/screens/artist/:sonyArtistId/countries?startDate=:startDate&endDate=:endDate',
    'country': '/screens/artist/:sonyArtistId/country/:countryCode?startDate=:startDate&endDate=:endDate',
    'dashboard': '/screens/artist/:sonyArtistId/dashboard?startDate=:startDate&endDate=:endDate',
    'demographics': '/screens/artist/:sonyArtistId/demographics?startDate=:startDate&endDate=:endDate',
    'social': '/screens/artist/:sonyArtistId/social?startDate=:startDate&endDate=:endDate',
    'tracks': '/screens/artist/:sonyArtistId/tracks?startDate=:startDate&endDate=:endDate',
    'track': '/screens/track/:trackId?startDate=:startDate&endDate=:endDate'
}


def auth():
    conn = client.HTTPSConnection(os.environ.get('AUTH0_ISSUER').replace('https://', ''))
    client_id = os.environ.get('PERFORMANCE_TEST_CLIENT_ID')
    client_secret = os.environ.get('PERFORMANCE_TEST_CLIENT_SECRET')
    audience = os.environ.get('AUTH0_AUDIENCE')
    payload = '''
      {{
        "client_id": "{client_id}",
        "client_secret": "{client_secret}",
        "audience": "{audience}",
        "grant_type": "client_credentials"
      }}
    '''.format(client_id=client_id, client_secret=client_secret, audience=audience)
    headers = { 'content-type': "application/json" }
    conn.request("POST", "/oauth/token", payload, headers)
    data = json.loads(conn.getresponse().read().decode('utf-8'))
    try:
      return data["access_token"]
    except KeyError:
       raise KeyError("Authorization failed")


def get_lines(file_name):
    with open(file_name, 'r') as txt:
        return txt.read().splitlines()


def with_parameters(url, artist_id, start_date, end_date, country_code, track_id):
    return url \
    .replace(':sonyArtistId', artist_id) \
    .replace(':startDate', start_date) \
    .replace(':endDate', end_date) \
    .replace(':countryCode', country_code) \
    .replace(':trackId', track_id)


def get_url(url, artist_and_tracks, country_codes):
    country_code = random.choice(country_codes)
    artist_id, *tracks = artist_and_tracks.split()
    track_id = random.choice(tracks)
    end_date = date.today().isoformat()
    start_date = (date.today() - timedelta(days=28)).isoformat()
    return with_parameters(url, artist_id, start_date, end_date, country_code, track_id)


def try_connection():
    conn = client.HTTPSConnection(os.environ.get('TARGET_URL').replace('https://', ''))
    conn.request("GET", '/screens/artist/GRAS_7513', headers={ 'Authorization': 'Bearer ' + access_token })
    response = conn.getresponse().read().decode('utf-8')
    assert(response != 'Authentication Error')

access_token = auth()
try_connection()

country_codes = get_lines('./country_codes.txt')
artists_and_tracks = get_lines('./artists_and_tracks.txt')


def url_for(screen, artist_and_tracks):
    return get_url(urls[screen], artist_and_tracks, country_codes)


class Tasks(TaskSet):
    def get_screen(self, screen, artist_and_tracks):
        return self.client.get(url_for(screen, artist_and_tracks), headers={'Authorization':  'Bearer ' + access_token})

    # https://data-analytics.atlassian.net/browse/AMA-860
    @task
    def scenario_1(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)

    @task
    def scenario_2(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('tracks', artist_and_tracks)

    @task
    def scenario_3(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('tracks', artist_and_tracks)
        self.get_screen('track', artist_and_tracks)

    @task
    def scenario_4(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('social', artist_and_tracks)

    @task
    def scenario_5(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('artist', artist_and_tracks)

    @task
    def scenario_6(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('countries', artist_and_tracks)

    @task
    def scenario_7(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('countries', artist_and_tracks)
        self.get_screen('country', artist_and_tracks)

    @task
    def scenario_7(self):
        artist_and_tracks = random.choice(artists_and_tracks)
        self.get_screen('dashboard', artist_and_tracks)
        self.get_screen('country', artist_and_tracks)


class LoadTest(HttpLocust):
    host = os.environ.get('TARGET_URL')
    task_set = Tasks
    min_wait = 5000
    max_wait = 10000