"""Snowflake connection adapter."""

from functools import cache
from os import environ

import snowflake.connector as sc
from dotenv import load_dotenv

with open("get_results.sql", "r") as f:
    GET_RESULTS_SQL = f.read()

with open("show_create_table.sql", "r") as f:
    SHOW_CREATE_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(sc.DictCursor)
    return cursor.execute(sql_query)


if __name__ == "__main__":
    load_dotenv()
    TABLE_NAME = environ.get("TABLE_NAME") or None
    if not TABLE_NAME:
        raise RuntimeError("TABLE_NAME environment variable is not set.")
    create_table_results = query(SHOW_CREATE_TABLE_SQL.format(table_name=TABLE_NAME))
    print('CREATE TABLE', create_table_results.rowcount)
    for row in create_table_results:
        print(row)
    results = query(GET_RESULTS_SQL.format(table_name=TABLE_NAME))
    print('RESULTS', results.rowcount)
    print('first result', results.fetchone())
    # for row in results:
    #     print(row)