import time
import jwt
import json
import requests
import dateutil.parser as date_parser
from load.projects.cloud.tank_client.exceptions import JWTError
from load.projects.cloud.cloud_helper.utils import CachedToken


def get_jwt(iam_token_service_url, config):
    service_account_id = config['service_account_id']
    key_id = config['sa_key_id']
    with open(config['sa_private_key']) as f:
        private_key = f.read()

    now = int(time.time())
    payload = {
        'aud': iam_token_service_url,
        'iss': service_account_id,
        'iat': now,
        'exp': now + 360}

    # Формирование JWT.
    return jwt.encode(
        payload,
        private_key,
        algorithm='PS256',
        headers={'kid': key_id})


class JWTToken(CachedToken):

    @staticmethod
    def _get_with_no_cache(iam_token_service_url, config):
        try:
            data = {
                "jwt": get_jwt(iam_token_service_url, config).decode('utf-8')
            }
            with requests.Session() as session:
                resp = session.post(iam_token_service_url, data=json.dumps(data),
                                    headers={'Content-Type': 'application/json'}, verify=False)
                resp.raise_for_status()
                response = resp.json()
                expire_at = date_parser.parse(response["expiresAt"]).timestamp()
                return response["iamToken"], expire_at
        except Exception as e:
            raise JWTError(f"Couldn't get iam token from jwt: {e}.") from e
