"""DB Connector for retrieving Snowflake query metadata."""

from datetime import datetime, timedelta
import time
import urllib

from cached_property import cached_property
from cryptography.hazmat.primitives.serialization import load_der_private_key
from cryptography.hazmat.backends import default_backend
import jwt
import requests
from snowflake.connector.auth.keypair import AuthByKeyPair


class AuthByKeyPairCustom(AuthByKeyPair):
    """
        Key pair based authentication
    """
    LIFETIME = timedelta(seconds=1200)

    def __init__(self, private_key):
        """
        :param private_key: a byte array of der formats of private key
        """
        super().__init__(private_key)
        self._private_key = private_key
        self._jwt_token = ''
        self._jwt_token_exp = 0

    def get_jwt_token(self, account, user):
        """Get JWT token to use in further requests."""
        account = account.upper()
        user = user.upper()

        now = datetime.utcnow()

        private_key = load_der_private_key(
            data=self._private_key, password=None, backend=default_backend())
        public_key_fp = self.calculate_public_key_fingerprint(private_key)

        self._jwt_token_exp = now + self.LIFETIME
        payload = {
            self.ISSUER: "{}.{}.{}".format(account, user, public_key_fp),
            self.SUBJECT: "{}.{}".format(account, user),
            self.ISSUE_TIME: now,
            self.EXPIRE_TIME: self._jwt_token_exp
        }

        return jwt.encode(
            payload, private_key, algorithm=self.ALGORITHM).decode('utf-8')


class SnowflakeBadResponse(Exception):
    """Snowflake Bad Response.

    Raised in case of bad response from Snowflake server
    """

    def __init__(self, response=None):
        """Create an instance of Snowflake Bad Response exception.

        Args:
            response (requests.Response): Bad response from Snowflake
        """
        super().__init__(
            'Bad response from Snowflake server. Got {}'.format(
                response.status_code or 'unknown error'))
        self.response = response


class SnowflakeMetadataConnector:
    """Snowflake Metadata connector.

    Provides a way to check Snowflake query metadata
    """

    def __init__(self, sf_config):
        """Create a Snowflake API connector.

        Args:
            sf_config (dict): Dictionary of credentials.
        """
        self.sf_config = sf_config
        self.session = requests.Session()
        self.session.headers.update({
            'Accept': 'application/json',
            'Content-Type': 'application/json'})

    def authenticate(self):
        """Authenticate connector.

        Raises:
            SnowflakeAPIError: Invalid credentials.
        """
        self.session.headers.update({
            'Authorization': 'Snowflake Token="{}"'.format(self._token)})

    def _get_url(self, path):
        return urllib.parse.urljoin(self._host, path)

    @cached_property
    def _host(self):
        return 'https://{}.snowflakecomputing.com'.format(
            self.sf_config['account'])

    @cached_property
    def _credentials(self):
        request_payload = {
            'ACCOUNT_NAME': self.sf_config['account'],
            'LOGIN_NAME': self.sf_config['user'],
            'PASSWORD': self.sf_config['password']
        }

        if self.sf_config.get('private_key'):
            auth_instance = AuthByKeyPairCustom(
                self.sf_config.get('private_key'))
            jwt_token = auth_instance.get_jwt_token(
                self.sf_config['account'], self.sf_config['user'])

            request_payload.update({
                'AUTHENTICATOR': 'SNOWFLAKE_JWT',
                'TOKEN': jwt_token})

        response = self.session.post(
            self._get_url('/session/v1/login-request'),
            json={'data': request_payload})
        if not response.ok:
            raise SnowflakeBadResponse(response)
        return response.json()['data']

    @property
    def _token(self):
        return self._credentials['token']

    def _get_query_data(self, snowflake_query_id):
        response = self.session.get(
            self._get_url('/monitoring/queries/{}'.format(snowflake_query_id)))
        if not response.ok:
            time.sleep(5)  # sleep and retry once
            response = self.session.get(
                self._get_url(
                    '/monitoring/queries/{}'.format(snowflake_query_id)))
            if not response.ok:
                raise SnowflakeBadResponse(response)
        return response.json()['data']

    def get_query_scan_bytes_number(self, cursor):
        """Get number of bytes scanned not from cache.

        Args:
            cursor (snowflake.connector.cursor.SnowflakeCursor):
                Snowflake query cursor

        Returns:
            int: number of bytes scanned not from cache
        """
        data = self._get_query_data(cursor.sfqid)
        return sum(
            query['stats'].get('ioRemoteFdnReadBytes', 0)
            for query in data['queries'])

    def get_query_stats_by_sfqid(self, sfqid):
        """Get number of bytes scanned not from cache.

        Args:
            sfqid (str): Snowflake query id.

        Returns:
            dict: number of bytes scanned from cache, number of bytes scanned
                from remote storage.
        """
        remote_bytes = 0
        local_bytes = 0

        data = self._get_query_data(sfqid)
        if data:
            remote_bytes = sum(
                query.get('stats').get('ioRemoteFdnReadBytes', 0)
                for query in data['queries'])
            local_bytes = sum(
                query.get('stats').get('ioLocalFdnReadBytes', 0)
                for query in data['queries'])

        return {
            'remote_bytes': remote_bytes,
            'local_bytes': local_bytes
        }