from dagster import EnvVar
from dagster_snowflake import SnowflakeResource
from src.utils.database import get_snowflake_private_key
from src.utils.snowflake.constants import Role, Warehouse

# If there is an EnvVar for private_key_path, use that. Otherwise, use private_key.
if EnvVar("SNOWFLAKE_PRIVATE_KEY_PATH").get_value():
    snowflake = SnowflakeResource(
        account=EnvVar("SNOWFLAKE_ACCOUNT"),
        user=EnvVar("SNOWFLAKE_ACCOUNT_USER"),
        private_key_path=EnvVar("SNOWFLAKE_PRIVATE_KEY_PATH"),
        database=EnvVar("SNOWFLAKE_DB"),
        schema=EnvVar("SNOWFLAKE_SCHEMA"),
        role=Role.ACCOUNTADMIN,
        warehouse=Warehouse.SMALL,  # Default warehouse, can be overridden at runtime
    )
else:
    snowflake = SnowflakeResource(
        account=EnvVar("SNOWFLAKE_ACCOUNT"),
        user=EnvVar("SNOWFLAKE_ACCOUNT_USER"),
        private_key=EnvVar("SNOWFLAKE_PRIVATE_KEY").get_value().replace("\\n", "\n"),
        database=EnvVar("SNOWFLAKE_DB"),
        schema=EnvVar("SNOWFLAKE_SCHEMA"),
        role=EnvVar("SNOWFLAKE_ACCOUNT_ROLE"),
        warehouse=Warehouse.SMALL,  # Default warehouse, can be overridden at runtime
    )

__all__ = ["snowflake"]