import sys

from locust import HttpUser
from locust import events

from auth import get_auth_class
from metrics import send_stats_to_datadog
from tasks import get_locust_tasks, get_ignored_http_statuses_and_urls


ignored_http_statuses_and_urls = get_ignored_http_statuses_and_urls()


class LoadtestUser(HttpUser):
    tasks = get_locust_tasks()
    auth_class = get_auth_class()

    def on_start(self) -> None:
        if not self.tasks:
            sys.exit(1)
        self.login()

    def login(self) -> None:
        if self.auth_class:
            self.client.headers.update({"Authorization": self.auth_class.get_token()})


@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    environment.events.request._handlers = []

    def on_request(response, request_type, name, response_time, response_length, exception=None, **_kwargs):
        if response.status_code == 404:
            return

        for status, url_pattern in ignored_http_statuses_and_urls:
            if response.status_code == int(status) and url_pattern in response.url:
                return

        environment.stats.log_request(request_type, name, response_time, response_length)
        if exception:
            environment.stats.log_error(request_type, name, exception)

    environment.events.request.add_listener(on_request)


@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    send_stats_to_datadog(environment)