"""
Snowflake connection management for test-runner.

Handles creation and configuration of Snowflake connections for different modes:
- Named connections (via SnowflakeConnectorFactory)
- Session token-based connections (for desktop app integration)
"""

from __future__ import annotations

import logging
from typing import Any

from snowflake.connector import SnowflakeConnection
from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.snowflake.connector.connector_factory_snowflake import (
    SnowflakeConnectorFactory,
)
from snowflake.snowflake_data_validation.snowflake.model.snowflake_named_connection import (
    SnowflakeNamedConnection,
)

from test_runner.validate.connector_wrapper import create_connector_from_connection
from test_runner.validate.targets.snowflake_connection_config import (
    ConnectionMode,
    SessionTokenConfig,
)

LOGGER = logging.getLogger(__name__)


class SnowflakeConnectionManager:
    """Manages Snowflake connection creation for different modes.
    
    Supports:
    - Named connections (from TOML files via snowflake-data-validation)
    - Session token connections (for connection sharing)
    - Connection wrapping (reuse existing connections)
    """

    def create_connector(self, config: SnowflakeNamedConnection | SessionTokenConfig) -> ConnectorBase:
        """Create a connector from configuration.
        
        Args:
            config: Either a SnowflakeNamedConnection or SessionTokenConfig
            
        Returns:
            ConnectorBase wrapping the Snowflake connection
        """
        if isinstance(config, SessionTokenConfig):
            return self._create_session_token_connector(config)
        
        return SnowflakeConnectorFactory.create_connector(config)

    def _create_session_token_connector(self, config: SessionTokenConfig) -> ConnectorBase:
        """Create connector using session tokens.
        
        Args:
            config: Validated session token configuration
            
        Returns:
            ConnectorBase wrapping the session token connection
        """
        import snowflake.connector
        
        LOGGER.debug("Creating session token connection for account=%s user=%s", 
                    config.account, config.user)
        
        connection = snowflake.connector.connect(**config.to_connection_params())
        
        return create_connector_from_connection(connection, owns_connection=True)

    def wrap_connection(
        self,
        connection: SnowflakeConnection,
        *,
        owns_connection: bool = False
    ) -> ConnectorBase:
        """Wrap an existing Snowflake connection.
        
        Args:
            connection: Active Snowflake connection
            owns_connection: If True, the connector will close the connection
                           when disposed. If False, caller manages lifecycle.
        
        Returns:
            ConnectorBase wrapping the connection
        """
        return create_connector_from_connection(
            connection,
            owns_connection=owns_connection
        )
