import base64
import os
import snowflake.connector
import toml

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization


def get_snowflake_connection_instance():
    config_path = os.path.expanduser("~/.crux/snowflake/config.toml")
    config = toml.load(config_path).get("us_account", {})

    key_b64 = os.environ.get("SNOWFLAKE_PRIVATE_KEY") or config.get("private_key")
    keypass = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PWD") or config.get(
        "private_key_password"
    )
    user = os.environ.get("SNOWFLAKE_USER_NAME") or config.get("user")
    account = os.environ.get("SNOWFLAKE_ACCOUNT_NAME") or config.get("account")
    warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE") or config.get("warehouse")
    database = os.environ.get("SNOWFLAKE_DATABASE_NAME") or config.get("database")
    schema = os.environ.get("SNOWFLAKE_SCHEMA_NAME") or config.get("schema")
    role = os.environ.get("SNOWFLAKE_USER_ROLE") or config.get("role")

    decoded_credential = base64.b64decode(key_b64)
    p_key = serialization.load_pem_private_key(
        decoded_credential,
        password=keypass.encode(),
        backend=default_backend(),
    )

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

    conn = snowflake.connector.connect(
        user=user,
        private_key=pkb,
        account=account,
        warehouse=warehouse,
        database=database,
        schema=schema,
        role=role,
    )
    return conn


def get_snowflake_eu_connection_instance():
    config_path = os.path.expanduser("~/.crux/snowflake/config.toml")
    config = toml.load(config_path).get("eu_account", {})

    key_b64 = os.environ.get("SNOWFLAKE_PRIVATE_KEY") or config.get("private_key")
    keypass = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PWD") or config.get(
        "private_key_password"
    )
    user = os.environ.get("SNOWFLAKE_USER_NAME") or config.get("user")
    account = os.environ.get("SNOWFLAKE_ACCOUNT_NAME") or config.get("account")
    warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE") or config.get("warehouse")
    database = os.environ.get("SNOWFLAKE_DATABASE_NAME") or config.get("database")
    schema = os.environ.get("SNOWFLAKE_SCHEMA_NAME") or config.get("schema")
    role = os.environ.get("SNOWFLAKE_USER_ROLE") or config.get("role")

    decoded_credential = base64.b64decode(key_b64)
    p_key = serialization.load_pem_private_key(
        decoded_credential,
        password=keypass.encode(),
        backend=default_backend(),
    )

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

    conn = snowflake.connector.connect(
        user=user,
        private_key=pkb,
        account=account,
        warehouse=warehouse,
        database=database,
        schema=schema,
        role=role,
    )
    return conn


def get_snowflake_vps_connection_instance():
    config_path = os.path.expanduser("~/.crux/snowflake/config.toml")
    config = toml.load(config_path).get("vps_account", {})

    key_b64 = os.environ.get("SNOWFLAKE_PRIVATE_KEY") or config.get("private_key")
    keypass = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PWD") or config.get(
        "private_key_password"
    )
    user = os.environ.get("SNOWFLAKE_USER_NAME") or config.get("user")
    account = os.environ.get("SNOWFLAKE_ACCOUNT_NAME") or config.get("account")
    warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE") or config.get("warehouse")
    database = os.environ.get("SNOWFLAKE_DATABASE_NAME") or config.get("database")
    schema = os.environ.get("SNOWFLAKE_SCHEMA_NAME") or config.get("schema")
    role = os.environ.get("SNOWFLAKE_USER_ROLE") or config.get("role")

    decoded_credential = base64.b64decode(key_b64)
    p_key = serialization.load_pem_private_key(
        decoded_credential,
        password=keypass.encode(),
        backend=default_backend(),
    )

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

    conn = snowflake.connector.connect(
        user=user,
        private_key=pkb,
        account=account,
        warehouse=warehouse,
        database=database,
        schema=schema,
        role=role,
    )
    return conn
