"""This file and its contents are licensed under the Apache License 2.0. Please see the included NOTICE for copyright information and LICENSE for a copy of the license.
"""
import json
import random
import string

from locust import HttpUser, TaskSet, between, task


def randomString(stringLength):
    """Generate a random string of fixed length"""
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(stringLength))


all_labels = [
    'Person',
    'Organization',
    'Fact',
    'Money',
    'Date',
    'Time',
    'Ordinal',
    'Percent',
    'Product',
    'Language',
    'Location',
]


def get_result(text):
    start = random.randint(0, len(text))
    end = min(len(text), start + random.randint(3, 30))
    results = []
    for i in range(random.randint(1, 10)):
        results.append(
            {
                'type': 'labels',
                'from_name': 'ner',
                'to_name': 'text',
                'value': {'labels': [random.choice(all_labels)], 'start': start, 'end': end},
            }
        )
    return results


class UserWorksWithProject(TaskSet):
    def on_start(self):
        r = self.client.get('/api/annotator/projects')
        all_projects = r.json()
        self.project_id = random.choice(all_projects)

    @task(100)
    def complete_task_via_api(self):
        r = self.client.get('/api/projects/%i/next/' % self.project_id, name='/api/projects/<id>/next')
        task = r.json()
        task_id = task['id']
        task_text = task['data']['text']
        results = get_result(task_text)
        payload = json.dumps({'result': results})
        headers = {'content-type': 'application/json'}
        self.client.post(
            '/api/tasks/%i/annotations' % task_id, payload, headers=headers, name='/api/tasks/<id>/annotations'
        )

    @task(1)
    def stop(self):
        self.interrupt()


class WebsiteUser(HttpUser):
    wait_time = between(3, 9)

    tasks = {UserWorksWithProject: 10}

    def on_start(self):
        self.login()

    def login(self):
        response = self.client.get('/')
        csrftoken = response.cookies['csrftoken']
        num_collabs = 100
        payload = {'email': f'collab_{random.randint(0, num_collabs)}@loadtests.me', 'password': '123456789'}
        self.client.post('/annotator/login', payload, headers={'X-CSRFToken': csrftoken})