"""Snowflake connection wrapper for the SDK.

Wraps ``src.utils.snowflake_config`` for connection management and
``src.persistence.snowflake_bulk.SnowflakeBulkLoader`` for bulk ops.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional


class SnowflakeConnection:
    """Managed Snowflake connection for SDK operations.

    Reuses:
        - ``src.utils.snowflake_config.read_snowflake_connection()``
        - ``src.utils.snowflake_config.connect_snowflake()``
        - ``src.persistence.snowflake_bulk.SnowflakeBulkLoader``

    Usage::

        conn = SnowflakeConnection(connection_name="my_conn")
        conn.connect()
        rows = conn.execute("SELECT 1")
        conn.close()
    """

    def __init__(
        self,
        connection_name: Optional[str] = None,
        config_override: Optional[Dict[str, Any]] = None,
    ):
        """Initialize connection wrapper.

        Args:
            connection_name: Named connection from ``~/.snowflake/config.toml``.
                If None, uses the default connection.
            config_override: Direct config dict (account, user, password, etc.)
                that bypasses TOML lookup.
        """
        self._connection_name = connection_name
        self._config_override = config_override
        self._conn = None
        self._config: Dict[str, Any] = {}

    @property
    def is_connected(self) -> bool:
        return self._conn is not None

    def connect(self) -> None:
        """Establish connection to Snowflake.

        Reuse: ``src.utils.snowflake_config.read_snowflake_connection()``
        then ``src.utils.snowflake_config.connect_snowflake()``.

        Raises:
            sdk.core.exceptions.ConnectionError: On connection failure.
        """
        from databridge_core.sdk.core.exceptions import ConnectionError as SDKConnectionError

        try:
            if self._config_override:
                self._config = self._config_override
            else:
                from src.utils.snowflake_config import read_snowflake_connection

                self._config = read_snowflake_connection(
                    connection_name=self._connection_name
                )

            from src.utils.snowflake_config import connect_snowflake

            self._conn = connect_snowflake(self._config)
        except Exception as exc:
            raise SDKConnectionError(
                f"Failed to connect to Snowflake: {exc}"
            ) from exc

    def execute(self, sql: str, params: Optional[Dict[str, Any]] = None) -> List[Any]:
        """Execute a SQL statement and return rows.

        Args:
            sql: SQL statement to execute.
            params: Optional bind parameters.

        Returns:
            List of result rows.

        Raises:
            sdk.core.exceptions.ConnectionError: If not connected.
        """
        from databridge_core.sdk.core.exceptions import ConnectionError as SDKConnectionError

        if not self._conn:
            raise SDKConnectionError("Not connected. Call connect() first.")
        cursor = self._conn.cursor()
        try:
            cursor.execute(sql, params or {})
            return cursor.fetchall()
        finally:
            cursor.close()

    def execute_ddl(self, ddl: str) -> str:
        """Execute DDL (CREATE, ALTER, DROP) and return status message.

        Args:
            ddl: DDL statement.

        Returns:
            Status message from Snowflake.
        """
        rows = self.execute(ddl)
        return rows[0][0] if rows else "OK"

    def close(self) -> None:
        """Close the connection."""
        if self._conn:
            self._conn.close()
            self._conn = None

    def __enter__(self) -> "SnowflakeConnection":
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()
