import os
import argparse


from snowflake.snowpark.session import Session
from dotenv import load_dotenv

# Initialize parser
parser = argparse.ArgumentParser(description="A script that accepts arguments.")

# Add arguments
parser.add_argument("--table", required=True)
parser.add_argument("--name", required=True)
parser.add_argument("--type", required=True)


args = parser.parse_args()

# Load environment variables
load_dotenv()
SNOWFLAKE_CONFIGS = {
    "account": os.getenv("SNOWFLAKE_ACCOUNT"),
    "user": os.getenv("SNOWFLAKE_ACCOUNT_USER"),
    "password": os.getenv("SNOWFLAKE_ACCOUNT_PASSWORD"),
    "role": os.getenv("SNOWFLAKE_ACCOUNT_ROLE"),
}

# Get Snowflake session
def get_snowflake_session(database: str, warehouse: str, schema: str) -> Session:
    return Session.builder.configs(
        {"warehouse": warehouse, "database": database, "schema": schema, **SNOWFLAKE_CONFIGS}
    ).create()


# cd to the same directory as this file and run the following script to deploy the procedure
# make sure you push the code to main branch before deploy
# uv run deploy_to_prod.py --file <file path>
if __name__ == "__main__":
    session = get_snowflake_session(database="SUNO_PROD", warehouse="SUNO_PROD_ENGINEER_X_SMALL", schema="PROD")

    try:
        # fetch the lates code from the git repository
        sql = '''
            alter table {table} add column {name} {type};
        '''.format(table=args.table, name=args.name, type=args.type)
        print (sql)
        print (session.sql(sql).collect())
    except Exception as e:
        print(e)
    finally:
        session.close()