from airflow.models import Variable
from airflow.providers.http.hooks.http import HttpHook
from requests.auth import AuthBase

from common.config.app import SERVICE, ENVIRONMENT


class AtlasAuth(AuthBase):
    atlas_conn_id = "atlas_http"

    headers = {
        "Client-Service": SERVICE,
        "Client-Environment": ENVIRONMENT
    }

    def __init__(self, _login: str, _password: str):
        pass

    def _get_token(self):
        hook = HttpHook(method="POST", http_conn_id=self.atlas_conn_id)
        response = hook.run(endpoint="oauth/token", headers=self.headers,
                            data=Variable.get("atlas_payload", deserialize_json=True))
        result = response.json()

        return f"{result['token_type']} {result['access_token']}"

    def __call__(self, r):
        r.headers.update(self.headers, **{"Authorization": self._get_token()})
        return r