"""
Snowflake connector wrapper that reuses an existing Snowflake connection.

This module provides a wrapper around the snowflake-data-validation
ConnectorSnowflake to enable connection sharing. Instead of creating
a new connection, it wraps an existing snowflake.connector.SnowflakeConnection.
"""

from __future__ import annotations

from typing import Any

import snowflake.connector
from snowflake.connector import SnowflakeConnection
from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase


class SnowflakeConnectorWrapper(ConnectorBase):
    """
    Wrapper around an existing Snowflake connection for data-validation framework.
    
    This allows reusing an established connection instead of creating a new one,
    which is essential for:
    - Session token-based connections (connection sharing from external systems)
    - Connection pooling
    - Avoiding redundant authentication
    - Maintaining consistent connection state
    
    The wrapper implements the ConnectorBase interface required by
    snowflake-data-validation while delegating to an existing connection.
    """

    def __init__(
        self,
        connection: SnowflakeConnection,
        *,
        owns_connection: bool = False
    ) -> None:
        """
        Initialize the wrapper with an existing Snowflake connection.
        
        Args:
            connection: An active snowflake.connector.SnowflakeConnection instance.
            owns_connection: If True, this wrapper will close the connection
                           when close() is called. If False (default), the
                           connection is assumed to be managed externally and
                           won't be closed by this wrapper.
        """
        self._connection = connection
        self._owns_connection = owns_connection
        self._closed = False

    @property
    def connection(self) -> SnowflakeConnection:
        """Return the wrapped Snowflake connection."""
        if self._closed:
            raise RuntimeError("Connection wrapper has been closed")
        return self._connection

    def connect(self) -> SnowflakeConnection:
        """Return the existing connection (required by ConnectorBase interface)."""
        return self.connection

    def execute_query(self, query: str) -> list[dict[str, Any]]:
        """
        Execute a SELECT query and return results as a list of dicts.
        
        Args:
            query: SQL SELECT statement to execute.
            
        Returns:
            List of dictionaries where keys are column names.
        """
        cursor = self._connection.cursor(snowflake.connector.DictCursor)
        try:
            cursor.execute(query)
            return cursor.fetchall()
        finally:
            cursor.close()

    def execute_statement(self, statement: str) -> None:
        """
        Execute a DDL/DML statement without returning results.
        
        Args:
            statement: SQL statement to execute (CREATE, INSERT, etc.).
        """
        cursor = self._connection.cursor()
        try:
            cursor.execute(statement)
        finally:
            cursor.close()

    def execute_query_no_return(self, query: str) -> None:
        """
        Execute a query without returning results (required by ConnectorBase interface).
        
        Args:
            query: SQL statement to execute.
        """
        self.execute_statement(query)

    def close(self) -> None:
        """
        Close the connector (and optionally the underlying connection).
        
        If owns_connection=True, closes the underlying Snowflake connection.
        If owns_connection=False, marks the wrapper as closed but leaves
        the connection open for external management.
        """
        if self._closed:
            return
            
        self._closed = True
        
        if self._owns_connection and self._connection:
            try:
                self._connection.close()
            except Exception:
                pass  # Best effort close

    def __enter__(self) -> SnowflakeConnectorWrapper:
        """Support context manager protocol."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Support context manager protocol."""
        self.close()


def create_connector_from_connection(
    connection: SnowflakeConnection,
    *,
    owns_connection: bool = False
) -> ConnectorBase:
    """
    Factory function to create a ConnectorBase from an existing connection.
    
    This is the primary entry point for creating a connector wrapper when
    you have an existing Snowflake connection that you want to reuse.
    
    Args:
        connection: Active Snowflake connection.
        owns_connection: Whether the connector should close the connection
                        when it's closed.
    
    Returns:
        ConnectorBase instance that wraps the provided connection.
    
    Example:
        >>> import snowflake.connector
        >>> conn = snowflake.connector.connect(...)
        >>> connector = create_connector_from_connection(conn)
        >>> # Use connector with data-validation framework
        >>> connector.close()  # Connection remains open
    """
    return SnowflakeConnectorWrapper(connection, owns_connection=owns_connection)
