import os
import logging
import dataiku
from dataiku.snowpark import DkuSnowpark

logger = logging.getLogger("snowflakecortexplugin")

class SnowflakeConnectionHelper:
    def __init__(self,snowflake_connection_name: str):
        self.snowflake_connection_name = snowflake_connection_name

        # Get project and connection info
        project_key = dataiku.default_project_key()
        client = dataiku.api_client()
        connection = client.get_connection(self.snowflake_connection_name)
        connection_info = connection.get_info(contextual_project_key=project_key)

        # Get account info from Snowpark (read-only, doesn't consume token)
        dku_snowpark = DkuSnowpark()
        self.snowpark_session = dku_snowpark.get_session(
            connection_name=self.snowflake_connection_name,
            project_key=project_key
        )

        # Handle proxy if configured
        jdbc_props = connection_info.get('params', {}).get('properties', [])
        if jdbc_props:
            props_dict = {p["name"]: p["value"] for p in jdbc_props if "name" in p and "value" in p}
            if all(k in props_dict for k in ["useProxy", "proxyHost", "proxyPort"]):
                if props_dict["useProxy"].lower() in ("on", "true"):
                    proxy_url = f"http://{props_dict['proxyHost']}:{props_dict['proxyPort']}"
                    for var in ['HTTP_PROXY', 'HTTPS_PROXY', 'http_proxy', 'https_proxy']:
                        os.environ[var] = proxy_url
                    logger.info(f"Configured proxy: {proxy_url}")

        self.snowflake_host = connection_info.get('resolvedParams', {}).get('host')

        if not self.snowflake_host:
            raise ValueError("Snowflake host not found in connection.")

        # Get API auth token
        token = connection_info.get("resolvedSnowflakeAuthToken",{}).get("bearerToken")
        token_type = connection_info.get('resolvedSnowflakeAuthToken', {}).get("snowflakeAuthorizationTokenType")

        if not token or not token_type:
            raise ValueError("Snowflake Auth token not resolved in connection.")

        self.headers = {
            "Authorization": f"Bearer {token}",
            "X-Snowflake-Authorization-Token-Type": token_type,
            "Content-Type": "application/json",
        }

        logger.info(f"Resolved Snowflake connection settings")
