import base64
import json
import os

from airflow.providers.amazon.aws.hooks.secrets_manager import SecretsManagerHook

from lib.config import SECRETS_PREFIX
from lib.config import SF_ACCOUNT
from lib.config import SF_DATABASE
from lib.config import SF_ROLE
from lib.config import SF_SCHEMA
from lib.config import SF_USER
from lib.config import SF_WAREHOUSE


def set_connections():
    # Set connection to AWS (Localstack)
    os.environ['AIRFLOW_CONN_AWS_DEFAULT'] = json.dumps({
        'conn_type': 'aws',
        'login': 'localstack',
        'password': 'localstack',
        'extra': {
            'region_name': 'us-east-1',
            'endpoint_url': 'http://localstack:4566'
        }
    })

    # Set connection to Snowflake
    sm_hook = SecretsManagerHook()
    sf_private_key = sm_hook.get_secret(f'{SECRETS_PREFIX}/SF_PRIVATE_KEY')
    sf_passphrase = sm_hook.get_secret(f'{SECRETS_PREFIX}/SF_PASSPHRASE')
    sf_private_key_decoded = base64.b64decode(sf_private_key).decode()

    os.environ['AIRFLOW_CONN_SNOWFLAKE_DEFAULT'] = json.dumps({
        'conn_type': 'snowflake',
        'login': SF_USER,
        'password': sf_passphrase,
        'schema': SF_SCHEMA,
        'extra': {
            'account': SF_ACCOUNT,
            'database': SF_DATABASE,
            'warehouse': SF_WAREHOUSE,
            'role': SF_ROLE,
            'private_key_content': sf_private_key_decoded
        }
    })