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

try:
    SNOWFLAKE_PRIVATE_KEY_PATH = os.environ["SNOWFLAKE_PRIVATE_KEY_PATH"]
    SNOWFLAKE_KEY_PASSPHRASE = os.environ["SNOWFLAKE_KEY_PASSPHRASE"]

    with open(SNOWFLAKE_PRIVATE_KEY_PATH, "rb") as key:
        p_key = serialization.load_pem_private_key(
            key.read(),
            # if you have no passphrase, set
            # password=None,
            # here and comment out the line below
            password=SNOWFLAKE_KEY_PASSPHRASE.encode(),
            backend=default_backend(),
        )
    SNOWFLAKE_PRIVATE_KEY = p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
except Exception:
    raise ValueError(
        "SNOWFLAKE_PRIVATE_KEY_PATH and SNOWFLAKE_KEY_PASSPHRASE must be set in the environment variables"
    )


def reset_snowflake_schema():
    # Get environment variables
    database = os.environ["SNOWFLAKE_DATABASE"]
    schema = os.environ["SNOWFLAKE_SCHEMA"]
    if schema.upper() == "PROD":
        raise ValueError("Cannot reset the PROD schema")
    if not database or not schema:
        raise ValueError("SNOWFLAKE_DATABASE and SNOWFLAKE_SCHEMA must be set in the environment variables")

    # Connect to Snowflake
    conn = snowflake.connector.connect(
        user=os.getenv("SNOWFLAKE_USER"),
        account="sme-delphi",
        warehouse=os.getenv("SNOWFLAKE_WAREHOUSE"),
        database=database,
        schema=schema,
        role=os.getenv("SNOWFLAKE_ROLE"),
        private_key=SNOWFLAKE_PRIVATE_KEY
    )

    try:
        cur = conn.cursor()
        # Drop the schema if it exists
        cur.execute(f"DROP SCHEMA IF EXISTS {database}.{schema}")
        # Create the schema
        cur.execute(f"CREATE SCHEMA {database}.{schema}")
        print(f"Schema {schema} in database {database} has been reset.")
    finally:
        # Close the cursor and connection
        cur.close()
        conn.close()


# Example usage
if __name__ == "__main__":
    reset_snowflake_schema()