"""Proxy for snowflake requests."""
import logging
import time
from contextlib import contextmanager

from snowflake.connector.errors import ProgrammingError

from snowflake_connector import snowflake_conn

logger = logging.getLogger(__name__)

DEFAULT_CONFIG = {
    'pool_pre_ping': True,
    'pool_reset_on_return': 'rollback',
    'commit_before_close': False,
}

_LOCK_TIMEOUT_SECONDS = 30
_STATEMENT_TIMEOUT_SECONDS = 30
_MAX_RETRIES = 3
_RETRY_BACKOFF_SECONDS = 2
_SF_LOCK_TIMEOUT_SQLSTATE = '57014'


def fetchone(*args, **kwargs):
    """Passthrough to Snowflake to fetchone."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchone(*args, **kwargs)


def fetchall(*args, **kwargs):
    """Passthrough to Snowflake to fetchall."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchall(*args, **kwargs)


def execute(sql):
    """Passthrough to Snowflake to execute with session."""
    for attempt in range(1, _MAX_RETRIES + 1):
        with snowflake_conn.get_session() as session:
            try:
                session.execute(snowflake_conn.text(
                    f'ALTER SESSION SET LOCK_TIMEOUT = {_LOCK_TIMEOUT_SECONDS}'
                ))
                session.execute(snowflake_conn.text(
                    'ALTER SESSION SET ABORT_DETACHED_QUERY = TRUE'))
                session.execute(snowflake_conn.text(sql))
                session.commit()
                return
            except ProgrammingError as e:
                session.rollback()
                if _SF_LOCK_TIMEOUT_SQLSTATE in str(e) and attempt < _MAX_RETRIES:
                    logger.warning(
                        'SF lock contention on attempt %d/%d, retrying in %ds: %s',
                        attempt, _MAX_RETRIES, _RETRY_BACKOFF_SECONDS * attempt, e,
                    )
                    time.sleep(_RETRY_BACKOFF_SECONDS * attempt)
                else:
                    raise
            except Exception:
                session.rollback()
                raise


class _Transaction:
    """Handle for issuing statements inside an open Snowflake transaction."""

    def __init__(self, session):
        self._session = session

    def execute(self, sql, params=None):
        """Run one parametrized statement (`:name` binds) in the transaction.

        Args:
            sql (str): SQL statement with optional `:name` bind parameters.
            params (dict): values bound to the statement.
        """
        return self._session.execute(snowflake_conn.text(sql), params or {})

    def fetchall(self, sql, params=None):
        """Run a query and return all rows, within the transaction.

        Lets a read participate in the same transaction as the writes (e.g.
        reading current rules before an atomic create/delete).
        """
        return list(self.execute(sql, params).cursor.fetchall())

    def fetchone(self, sql, params=None):
        """Run a query and return the first row, within the transaction."""
        return self.execute(sql, params).cursor.fetchone()


@contextmanager
def transaction():
    """Run Snowflake reads and writes as one atomic transaction.

    Statements issued via the yielded handle share a single session and commit
    together when the block exits cleanly; any exception rolls the whole
    transaction back. Use for flows that must be all-or-nothing, e.g. reading
    current rules then performing the archive-then-delete soft delete.

    Yields:
        _Transaction: handle exposing ``execute``/``fetchall``/``fetchone``.
    """
    with snowflake_conn.get_session(commit_before_close=True) as session:
        session.execute(snowflake_conn.text(
            f'ALTER SESSION SET LOCK_TIMEOUT = {_LOCK_TIMEOUT_SECONDS}'))
        session.execute(snowflake_conn.text(
            'ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS'
            f' = {_STATEMENT_TIMEOUT_SECONDS}'))
        session.execute(snowflake_conn.text(
            'ALTER SESSION SET ABORT_DETACHED_QUERY = TRUE'))
        yield _Transaction(session)


def SQLLoader(path):
    """Create a SQL loader in given path."""
    return snowflake_conn.SQLLoader(path)