"""Load tests for ES."""

import os
import random
import sys
import time

from elasticsearch import Elasticsearch
from locust import HttpUser, TaskSet, task, events

_location = os.path.realpath(
    os.path.join(os.getcwd(), os.path.dirname(__file__)))

with open(os.path.join(_location, 'gp_names.txt')) as f:
    GP_NAMES = [name.strip() for name in f.readlines()]

with open(os.path.join(_location, 'gsr_names.txt')) as f:
    GSR_NAMES = [name.strip() for name in f.readlines()]


class Reader(TaskSet):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.search_engine = Elasticsearch(
            [os.environ.get('ELASTICSEARCH_URL')],
            http_auth=(
                os.environ.get('ELASTICSEARCH_LOGIN'),
                os.environ.get('ELASTICSEARCH_PASSWORD')))

    def _query(
            self, index='', body=None, query_name='', request_cache=False,
            request_type='POST'):
        start_time = time.perf_counter()
        try:
            result = self.search_engine.search(
                request_cache=request_cache,
                index=index,
                body=body
            )
        except Exception as e:
            total_time = int((time.perf_counter() - start_time) * 1000)
            events.request_failure.fire(
                request_type=request_type,
                name=query_name,
                response_time=total_time,
                response_length=0,
                exception=e)
        else:
            total_time = int((time.perf_counter() - start_time) * 1000)
            events.request_success.fire(
                request_type=request_type,
                name=query_name,
                response_time=total_time,
                response_length=sys.getsizeof(str(result)))

    @task(5)
    def globalparticipants_query_by_name(self):
        self._query(
            index='globalparticipants.v01',
            body={
                    "query": {
                        "match": {
                            "name": {
                                "query": random.choice(GP_NAMES),
                                "fuzziness": "AUTO"
                            }
                        }
                    }
                },
            query_name='Search globalparticipants.v01 by GP name',
            request_cache=False,
            request_type='POST'
        )

    @task(5)
    def globalsoundecordings_get_by_name(self):
        self._query(
            index='globalsoundrecordings.v01',
            body={
                    "query": {
                        "match": {
                            "name": {
                                "query": random.choice(GSR_NAMES),
                                "fuzziness": "AUTO"
                            }
                        }
                    }
                },
            query_name='Search globalsoundrecordings.v01 by GSR name',
            request_cache=False,
            request_type='POST'
        )


class LocustLoadOutput(HttpUser):
    tasks = {Reader: 1}
    min_wait = 2000
    max_wait = 30000