import os
import random
import string
from uuid import uuid4

import numpy as np
import pandas as pd
from locust import HttpUser, between, constant, events, tag, task


def get_project_id(client):
    with client.get('/api/projects', catch_response=True) as r:
        if r.status_code != 200:
            print(r.status_code)
            r.failure(r.status_code)
        else:
            project_list = r.json()
            print(project_list)
            ids = [p['id'] for p in project_list]
            if not ids:
                return
            return random.choice(ids)


def signup(client):
    username = str(uuid4())[:8]
    response = client.get('/')
    csrftoken = response.cookies['csrftoken']
    r = client.post(
        '/user/signup',
        {'email': f'{username}@heartex.com', 'password': 'password'},
        headers={'X-CSRFToken': csrftoken},
        catch_response=True,
    )
    print(f'User {username} signup response: {r.status_code}')
    return username


class Admin(HttpUser):
    weight = 1
    wait_time = constant(10)

    def on_start(self):
        username = signup(self.client)
        with self.client.post(
            '/api/projects',
            json={
                'title': f"{username}'s project",
                'label_config': '<View><Text name="text" value="$text"/><Choices name="label" toName="text"><Choice value="1"/><Choice value="2"/></Choices></View>',
            },
            catch_response=True,
        ) as r:
            if r.status_code != 201:
                r.failure(r.status_code)
            rdata = r.json()
            print('Get response: ', rdata)
            self.project_id = rdata['id']
            print(f'Project {self.project_id} has been created by user {self.client}')
            self.import_data()

    @task
    def view_data_manager(self):
        self.client.get(f'/projects/{self.project_id}/data', name='projects/<pk>/data')

    # @tag('import')
    # @task
    def import_data(self):
        self.client.post(
            '/api/projects/%i/import' % self.project_id,
            name='/api/projects/<pk>/import',
            files={'csv': open('data.csv', 'rb')},
            # headers={'content-type': 'multipart/form-data'})
        )


class Annotator(HttpUser):
    weight = int(os.environ.get('LOCUST_USERS')) - 1
    wait_time = between(1, 3)

    def on_start(self):
        signup(self.client)

    @tag('select project')
    @task(1)
    def select_project(self):
        self.project_id = get_project_id(self.client)

    @tag('labeling')
    @task(10)
    def do_labeling(self):
        if not hasattr(self, 'project_id') or not self.project_id:
            print('No projects yet...')
            return
        with self.client.get(
            f'/api/projects/{self.project_id}/next', name='/api/projects/<pk>/next', catch_response=True
        ) as r:
            task = r.json()
            self.client.post(
                f'/api/tasks/{task["id"]}/annotations',
                name='/api/tasks/<pk>/annotations',
                json={
                    'result': [
                        {
                            'from_name': 'label',
                            'to_name': 'text',
                            'type': 'choices',
                            'value': {'choices': [random.choice(['1', '2'])]},
                        }
                    ]
                },
            )


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


@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    rows = int(os.environ.get('IMPORTED_TASKS', 50000))
    print(f'Generating file with {rows} rows...')
    numbers = np.random.randint(low=0, high=100, size=(rows, 10)).tolist()
    columns = ['text'] + [f'col_{i}' for i in range(9)]
    pd.DataFrame(numbers, columns=columns).to_csv('data.csv')