from typing import Generator

import httpx


class UserAuth(httpx.Auth):
    requires_response_body = True

    def __init__(self, url: str, username: str, password: str):
        self.url = url
        self.basic_auth = httpx.BasicAuth(username, password)
        self.access_token = None

    def auth_flow(self, request) -> Generator[httpx.Request, httpx.Response, None]:
        if self.access_token is None:
            response = yield self.get_auth_token()
            response.raise_for_status()
            self.access_token = response.json()["data"]["v3_user_token"]
        request.headers["Authorization"] = f"Bearer {self.access_token}"
        yield request

    def get_auth_token(self) -> httpx.Request:
        request = httpx.Request(
            url=f"{self.url}/api/v3/auth/jwt?useBody=true",
            method="GET",
        )
        return next(self.basic_auth.auth_flow(request))


class APIClientAuth(httpx.Auth):
    requires_response_body = True

    def __init__(self, url: str, api_client_id: str, api_client_secret: str):
        self.url = url
        self.basic_auth = httpx.BasicAuth(api_client_id, api_client_secret)
        self.access_token = None

    def auth_flow(self, request) -> Generator[httpx.Request, httpx.Response, None]:
        if self.access_token is None:
            response = yield self.get_auth_token()
            response.raise_for_status()
            self.access_token = response.json()["access_token"]
        request.headers["Authorization"] = f"Bearer {self.access_token}"
        yield request

    def get_auth_token(self) -> httpx.Request:
        request = httpx.Request(
            url=f"{self.url}/api/v3/oauth/token?grant_type=client_credentials",
            method="POST",
        )
        return next(self.basic_auth.auth_flow(request))
