"""Snowflake connection adapter."""

from functools import cache
from os import environ

import snowflake.connector as sc
from dotenv import load_dotenv

with open("create_table.sql", "r") as f:
    CREATE_TABLE_SQL = f.read()

with open("load_from_stage.sql", "r") as f:
    LOAD_FROM_STAGE_SQL = f.read()

with open("create_file_format.sql", "r") as f:
    CREATE_FILE_FORMAT_SQL = f.read()

with open("create_proper_table.sql", "r") as f:
    CREATE_PROPER_TABLE_SQL = f.read()


@cache
def snowflake_connection():
    ctx = sc.connect(
        account=environ.get("SNOWFLAKE_ACCOUNT") or None,
        user=environ.get("SNOWFLAKE_USER_NAME") or None,
        authenticator="SNOWFLAKE_JWT",
        private_key_file=environ.get("SNOWFLAKE_PRIVATE_KEY_FILE") or None,
        private_key_file_pwd=environ.get("SNOWFLAKE_PRIVATE_KEY_FILE_PWD") or None,
        warehouse=environ.get("SNOWFLAKE_WAREHOUSE") or None,
        database=environ.get("SNOWFLAKE_DATABASE") or None,
        schema=environ.get("SNOWFLAKE_SCHEMA") or None,
    )
    return ctx


def query(sql_query):
    """Execute a query against Snowflake."""
    print('QUERY', sql_query)
    ctx = snowflake_connection()
    cursor = ctx.cursor()
    return cursor.execute(sql_query)


if __name__ == "__main__":
    load_dotenv()
    RUN_ID = environ.get("RUN_ID") or None
    if not RUN_ID:
        raise RuntimeError("RUN_ID environment variable is not set.")
    TABLE_NAME = environ.get("TABLE_NAME") or None
    if not TABLE_NAME:
        raise RuntimeError("TABLE_NAME environment variable is not set.")
    query(CREATE_FILE_FORMAT_SQL)
    # query(CREATE_TABLE_SQL.format(run_id=RUN_ID, table_name=TABLE_NAME))
    query(CREATE_PROPER_TABLE_SQL.format(run_id=RUN_ID, table_name=TABLE_NAME))
    query(LOAD_FROM_STAGE_SQL.format(run_id=RUN_ID, table_name=TABLE_NAME))