from datetime import datetime
from datetime import timedelta
from typing import Generator, Optional

import httpx
from pydantic import BaseModel
from pydantic import PrivateAttr


class AuthToken(BaseModel):
    token_type: str
    expires_in: int
    access_token: str
    _init_time: datetime = PrivateAttr(default_factory=datetime.utcnow)

    @property
    def expired(self):
        elapsed = datetime.utcnow() - self._init_time
        return elapsed > timedelta(seconds=self.expires_in)


class APIClientAuth(httpx.Auth):
    def __init__(self, url: str, api_client_id: str, api_client_secret: str):
        self.url = url
        self.api_client_id = api_client_id
        self.api_client_secret = api_client_secret
        self.token: Optional[AuthToken] = None

    def auth_flow(self, request) -> Generator[httpx.Request, httpx.Response, None]:
        request.headers["Authorization"] = f"Bearer {self.get_token()}"
        response = yield request
        if response.status_code == 401:
            request.headers["Authorization"] = f"Bearer {self.get_token()}"
            yield request

    def get_token(self) -> str:
        if self.token is None or self.token.expired:
            response = httpx.post(
                url=f"{self.url}/api/v3/oauth/token?grant_type=client_credentials",
                auth=(self.api_client_id, self.api_client_secret),
            )
            response.raise_for_status()
            self.token = AuthToken(**response.json())
        return self.token.access_token
