import time

from flask import Flask as App
from authlib.jose import JWT


class TestJwt:
    """
    Class to create JSON Web Tokens (JWTs) for testing purposes

    Supports generating tokens with a set of requested scopes (as either Azure 'roles' or 'scp's) using a testing
    signing key generated by the TestJwk class. Values for the `aud`, `iss` and `azp` claims in the payload come from
    the current Flask application.
    """

    _jwt = JWT()

    def __init__(self, *, app: App, roles: list = None, scps: list = None):
        """
        :type app: App
        :type app: Flask application
        :type roles: list
        :param scps: Optional scopes to include in the token (as a 'roles' claim) for testing authorisation
        :type roles: list
        :param scps: Optional scopes to include in the token (as a 'scp' claim) for testing authorisation
        """
        self.signing_key = app.config["TEST_JWKS"]

        self.header = {"alg": self.signing_key.algorithm, "kid": self.signing_key.kid()}
        self.payload = {
            "aud": app.config["AZURE_OAUTH_APPLICATION_ID"] or "testing",
            "exp": int(time.time() + 10000),
            "iat": int(time.time()),
            "iss": f"https://login.microsoftonline.com/{ app.config['AZURE_OAUTH_TENANCY'] or 'testing' }/v2.0",
            "nbf": int(time.time()),
            "sub": None,
            "azp": app.config["AZURE_OAUTH_CLIENT_APPLICATION_IDS"][0] or "testing",
        }
        self.scopes = set()
        if roles is not None:
            self.scopes.update(set(roles))
            self.payload["roles"] = " ".join(roles)
        if scps is not None:
            self.scopes.update(set(scps))
            self.payload["scp"] = " ".join(scps)

    def dumps(self) -> str:
        """
        Returns a signed/issued JWT encoded as a string for exchange

        :rtype str
        :return: Signed JWT
        """
        return self._jwt.encode(self.header, self.payload, self.signing_key.private_key_pem()).decode()
