"""Snowflake connection adapter."""

from functools import cache
from os import environ

import snowflake.connector as sc
from dotenv import load_dotenv

with open("create_final_table.sql", "r") as f:
    CREATE_FINAL_TABLE_SQL = f.read()

with open("show_create_table.sql", "r") as f:
    SHOW_CREATE_TABLE_SQL = f.read()

with open("pop_last_results.sql", "r") as f:
    POP_LAST_RESULTS_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_1 = environ.get("RUN_ID_1") or None
    # RUN_ID_2 = environ.get("RUN_ID_2") or None
    # run_schema_1 = query(SHOW_CREATE_TABLE_SQL.format(table_name=f'MODEL_REGRESSION_RESULTS_{RUN_ID_1}'))
    # for row in run_schema_1:
    #     print(row)
    # run_schema_2 = query(SHOW_CREATE_TABLE_SQL.format(table_name=f'MODEL_REGRESSION_RESULTS_{RUN_ID_2}'))
    # for row in run_schema_2:
    #     print(row)
    # query('ALTER TABLE ORCHARD_APP_REPORTING_V2.PROD_HIVE_MODEL_REPORTING.MODEL_REGRESSION_RESULTS_20260206T155452Z ADD COLUMN IF NOT EXISTS minimax VARCHAR(16777216);')
    # query('ALTER TABLE ORCHARD_APP_REPORTING_V2.PROD_HIVE_MODEL_REPORTING.MODEL_REGRESSION_RESULTS_20260206T155452Z ADD COLUMN IF NOT EXISTS mureka VARCHAR(16777216);')
    # query('ALTER TABLE ORCHARD_APP_REPORTING_V2.PROD_HIVE_MODEL_REPORTING.MODEL_REGRESSION_RESULTS_20260206T155452Z ADD COLUMN IF NOT EXISTS yue VARCHAR(16777216);')
    # query(POP_LAST_RESULTS_SQL)
    # query(CREATE_FINAL_TABLE_SQL)
    result = query('SELECT DISTINCT input_model, input_model_version FROM ORCHARD_APP_REPORTING_V2.PROD_HIVE_MODEL_REPORTING.HIVE_MODEL_RESPONSE_TRACKING;')
    for row in result:
        print(row)