"""Subclass of existing airflow snowflake hook."""

from functools import cached_property
from pathlib import Path

from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization

from lib import config


class RoyaltySnowflakeHook(SnowflakeHook):
    """Royalty flavored snowflake hook."""

    @cached_property
    def _get_conn_params(self):
        """Override SnowflakeHook._get_conn_params.

        In Dev, we authenticate with a private key file stored
        on the developer's machine.

        In QA, we authenticate with a private key whose content
        is stored in Secrets Manager.

        In Production, we authenticate with a password stored in
        Secrets Manager.

        TODO: Align the Production authentication method with QA to use
        a private key from Secrets Manager.
        """
        conn = self.get_connection(self.snowflake_conn_id)
        account = conn.extra_dejson.get('account', None)
        warehouse = conn.extra_dejson.get('warehouse', None)
        database = conn.extra_dejson.get('database', None)
        role = conn.extra_dejson.get('role', None)
        schema = conn.extra_dejson.get('schema')

        conn_config = {
            'user': conn.login,
            'password': conn.password,
            'schema': schema,
            'database': database,
            'account': account,
            'warehouse': warehouse,
            'role': role
        }

        private_key_pem = None
        passphrase = None

        if config.OWS_ENV == config.DEV_OWS_ENV:
            private_key_file = conn.extra_dejson.get('private_key_file')

            if private_key_file:
                private_key_file_path = Path(private_key_file)
                private_key_pem = Path(private_key_file_path).read_bytes()

            if conn.password:
                passphrase = conn.password.strip().encode()
        elif config.OWS_ENV == config.QA_OWS_ENV:
            if config.SNOWFLAKE_PRIVATE_KEY:
                private_key_pem = config.SNOWFLAKE_PRIVATE_KEY.encode()

            if config.SNOWFLAKE_KEY_PASSPHRASE:
                passphrase = config.SNOWFLAKE_KEY_PASSPHRASE.encode()

        if private_key_pem:
            p_key = serialization.load_pem_private_key(
                private_key_pem,
                password=passphrase,
                backend=default_backend()
            )

            pkb = p_key.private_bytes(
                encoding=serialization.Encoding.DER,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )

            conn_config['private_key'] = pkb
            conn_config.pop('password', None)

        return conn_config