from datetime import datetime, timezone
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:
        # We wrap it back into the standard PEM format
        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
    # .encode() turns the string into bytes for the library
    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",
        },
    )

objects_to_check = [
    {
        "rollup": "rollup_artist_summary",
        "luminate_fact": "VW_DAILY_FACT_ARTIST_SUMMARY_DS",
    },
    {
        "rollup": "rollup_mp_summary",
        "luminate_fact": "VW_DAILY_FACT_MP_SUMMARY_DS",
    },
    {
        "rollup": "rollup_mr_summary",
        "luminate_fact": "VW_DAILY_FACT_MR_SUMMARY_DS",
    },
    {
        "rollup": "rollup_mrel_summary",
        "luminate_fact": "VW_DAILY_FACT_MREL_SUMMARY_DS",
    },
    {
        "rollup": "rollup_mrelg_summary",
        "luminate_fact": "VW_DAILY_FACT_MRELG_SUMMARY_DS",
    },
    {
        "rollup": "rollup_song_summary",
        "luminate_fact": "VW_DAILY_FACT_SONG_SUMMARY_DS",
    },
]

is_luminate_views_updated_query = """
    select
        exists (
            select 1
            from luminate_db_listing_detail.extract_s.{luminate_fact}
            where modified_at > (select max(modified_at) from {database}.{schema}.{rollup})
        )
"""


def is_luminate_views_updated(conn):
    views_updated = 0
    for o in objects_to_check:
        q = is_luminate_views_updated_query.format(
            rollup=o["rollup"],
            luminate_fact=o["luminate_fact"],
            database=os.environ.get("SNOWFLAKE_DATABASE"),
            schema=os.environ.get("SNOWFLAKE_SCHEMA") + "_DBT_STAGE",
        )
        result = conn.cursor().execute(q).fetchone()
        if result[0] == True:
            print(f"{o['luminate_fact']} is updated")
            views_updated += 1
        elif result[0] == False:
            print(f"{o['luminate_fact']} is not updated")
            return False
        else:
            print("unknown result from is_luminate_views_updated:")
            print(result)
    return views_updated == len(objects_to_check)


def should_daily_run(conn):
    return is_luminate_views_updated(conn)


def should_monthly_run():
    return datetime.now(timezone.utc).day == 5

daily_run_options = [
    "--threads",
    "2",
]

daily_models = [
    "--select",
    "tag:common",
    "+tag:rtcw",
    "+tag:rtd",
    "+tag:wtd",
    "+tag:ytd",
]

monthly_run_options = [
    "--threads",
    "1",
]

monthly_models = [
    "--select",
    "tag:stage_monthly+",
]

def daily_run():
    conn = get_snowflake_connection()
    if should_daily_run(conn):
        daily_ls_args = ["dbt", "ls"] + daily_models
        print(" ".join(daily_ls_args))
        subprocess.run(daily_ls_args)
        daily_run_args = ["dbt", "run"] + daily_run_options + daily_models
        print(" ".join(daily_run_args))
        subprocess.run(daily_run_args, check=True)

        if should_monthly_run():
            monthly_ls_args = ["dbt", "ls"] + monthly_models
            print(" ".join(monthly_ls_args))
            subprocess.run(monthly_ls_args)
            monthly_run_args = ["dbt", "run"] + monthly_run_options + monthly_models
            print(" ".join(monthly_run_args))
            monthly_env = os.environ.copy()
            monthly_env["SNOWFLAKE_WAREHOUSE"] = os.environ.get("SNOWFLAKE_HEAVY_WAREHOUSE")
            subprocess.run(monthly_run_args, check=True, env=monthly_env)

if __name__ == '__main__':
    daily_run()