from dataclasses import dataclass
from datetime import datetime, timedelta

from airflow.models.variable import Variable
from airflow.providers.amazon.aws.hooks.secrets_manager import SecretsManagerHook

from dmp_workflows import config

OWS_JWT_DATA_VARIABLE_NAME = "ows_jwt"


@dataclass
class OwsJwt:
    token: str
    expires_at: datetime


class OwsJwtHook:
    _jwt_variable_name = "ows_jwt"

    @classmethod
    def _get_jwt_from_secrets_manager(cls) -> OwsJwt:
        hook = SecretsManagerHook(region_name=config.AWS_REGION)  # type: ignore[no-untyped-call]
        token = hook.get_secret(config.M2M_TOKEN_SECRET_KEY_NAME)
        token_expires_at = datetime.strptime(
            hook.get_secret(config.M2M_TOKEN_SECRET_EXPIRY_KEY_NAME),  # type: ignore[arg-type]
            "%Y-%m-%d %H:%M:%S.%f",
        )
        return OwsJwt(token, token_expires_at)  # type: ignore[arg-type]

    @classmethod
    def _get_jwt_from_cache(cls) -> OwsJwt | None:
        cached_jwt = Variable.get(
            cls._jwt_variable_name, deserialize_json=True, default_var=None
        )
        if not cached_jwt:
            return None
        return OwsJwt(
            token=cached_jwt["token"],
            expires_at=datetime.fromisoformat(cached_jwt["expires_at"]),
        )

    @classmethod
    def _save_jwt_to_cache(cls, jwt: OwsJwt) -> None:
        serializable_jwt = {
            "token": jwt.token,
            "expires_at": jwt.expires_at.isoformat(),
        }
        Variable.set(cls._jwt_variable_name, serializable_jwt, serialize_json=True)

    @classmethod
    def _has_jwt_expired(cls, jwt: OwsJwt) -> bool:
        return jwt.expires_at < (
            datetime.utcnow()
            + timedelta(
                seconds=60  # to ensure that token does not expire during the request
            )
        )

    @classmethod
    def get_jwt_token(cls) -> str:
        cached_jwt = cls._get_jwt_from_cache()
        if cached_jwt and not cls._has_jwt_expired(cached_jwt):
            return cached_jwt.token
        else:
            jwt = cls._get_jwt_from_secrets_manager()
            cls._save_jwt_to_cache(jwt)
            return jwt.token