import os
import subprocess
import snowflake.connector
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend

def get_snowflake_connection():
    # 1. Get the raw string and passphrase
    raw_key = os.environ.get("SNOWFLAKE_PRIVATE_KEY")
    passphrase = os.environ.get("SNOWFLAKE_KEY_PASSPHRASE")

    # 2. Reconstruct the PEM format if headers are missing
    if "BEGIN" not in raw_key:
        p_key_pem = (
            "-----BEGIN ENCRYPTED PRIVATE KEY-----\n"
            f"{raw_key}\n"
            "-----END ENCRYPTED PRIVATE KEY-----"
        )
    else:
        p_key_pem = raw_key

    # 3. Load and decrypt the key
    p_key_obj = serialization.load_pem_private_key(
        p_key_pem.encode(),
        password=passphrase.encode() if passphrase else None,
        backend=default_backend()
    )

    # 4. Convert to DER (binary) format as required by the connector
    pk_der = p_key_obj.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    )

    # 5. Connect
    return snowflake.connector.connect(
        user=os.environ.get("SNOWFLAKE_USER"),
        account="SME-DELPHI",
        authenticator="SNOWFLAKE_JWT",
        role=os.environ.get("SNOWFLAKE_ROLE"),
        warehouse=os.environ.get("SNOWFLAKE_LIGHT_WAREHOUSE"),
        database=os.environ.get("SNOWFLAKE_DATABASE"),
        schema=os.environ.get("SNOWFLAKE_SCHEMA"),
        private_key=pk_der,
        session_parameters={
            "TIMEZONE": "UTC",
        },
    )

staging_cmd = [
    "dbt",
    "run",
    "--select",
    "tag:staging",
]

mart_cmd = [
    "dbt",
    "run",
    "--select",
    "tag:mart",
]

mart_views_cmd = [
    "dbt",
    "run",
    "--select",
    "tag:mart_views",
]

def daily_run():
    print(" ".join(staging_cmd))
    subprocess.run(staging_cmd, check=True)
    print(" ".join(mart_cmd))
    subprocess.run(mart_cmd, check=True)
    print(" ".join(mart_views_cmd))
    subprocess.run(mart_views_cmd, check=True)

if __name__ == '__main__':
    daily_run()