import os

from locust import TaskSet as LocustTaskSet

from .auth import Authorization


def tag(*tags):
    def decorator_func(func):
        if not os.environ.get("TAG") or set(os.environ["TAG"].split(",")) & set(tags) & Tag.all():
            func.locust_task_weight = 1
        return func

    return decorator_func


class Tag:
    LOCAL = 'LOCAL'
    AGO = 'AGO'
    AP = 'AP'
    VENDOR = 'VENDOR'
    IMAGES = "IMAGES"
    NOTIFICATIONS = 'NOTIFICATIONS'
    ATLAS_UM = 'ATLAS_UM'

    @classmethod
    def all(cls) -> set:
        return set([getattr(cls, name) for name in dir(cls) if not name.startswith('_')])


class TaskSet(LocustTaskSet):
    auth_class = Authorization
    headers = {}
    api_prefix = ""
    validators = ()
    allow_redirects = True

    def on_start(self) -> None:
        if len(self.tasks) == 0:
            self.locust.stop_timeout = 0
            return

        self.login()
        self.client.headers.update(self.headers)

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

    def url(self, url_template: str, **kwargs) -> str:
        if self.api_prefix:
            url_template = f"{self.api_prefix}/{url_template}"

        return url_template.format(**kwargs)

    def request(self, method: str, **kwargs) -> None:
        if os.environ.get("DEBUG"):
            kwargs.pop("name", None)
        elif "name" not in kwargs:
            kwargs["name"] = kwargs["url"]

        if "allow_redirects" not in kwargs:
            kwargs["allow_redirects"] = self.allow_redirects

        with self.client.request(method, catch_response=True, **kwargs) as response:
            if response.status_code == 404 and not os.environ.get("ENV_CHECK"):
                response._is_reported = True
                return
            for validator in self.validators:
                validator(response)

    def get(self, **kwargs) -> None:
        self.request("GET", **kwargs)

    def post(self, **kwargs) -> None:
        self.request("POST", **kwargs)