"""DB Connector. Manages interactions with Snowflake."""

from contextlib import contextmanager
from sqlalchemy import text
from snowflake_connector.snowflake_conn import (
    _get_sessionmaker,
    validator,
    DBAPIError,
    DatabaseError,
    logger,
)


def fetchproxy(sql, params=None, sf_config=None, retry_on_disconnect=True):
    """Execute an SQL statement.

    Args:
        sql (str): parametrized sql statement.
        params (dict): A dict of params to bind to sql_template.
        sf_config (dict): Full or partial (without credentials) Snowflake
            config.
        retry_on_disconnect (bool): whether we should retry to execute
            the query on disconnect.

    """
    if params:
        sql, params = validator.format_identifiers(sql, params)

    while True:
        session = _get_sessionmaker(sf_config)()
        try:
            res = session.execute(text(sql), params)
            # Return mappings for SQLAlchemy 2.x dict-like row access
            return res.mappings()
        except DBAPIError as dpe:
            if not (retry_on_disconnect and isinstance(dpe.orig, DatabaseError)):
                raise
            logger.warning("Trying to re-execute due to database error: %s", str(dpe))
            retry_on_disconnect = False
        else:
            return res

        finally:
            session.close()


@contextmanager
def fetchproxy_stream(sql, params=None, sf_config=None, retry_on_disconnect=True):
    """Execute an SQL statement and yield the result cursor.

    Context manager that keeps the session open while the caller iterates.
    """
    if params:
        sql, params = validator.format_identifiers(sql, params)

    session = None
    try:
        while True:
            session = _get_sessionmaker(sf_config)()
            try:
                res = session.execute(text(sql), params)
                yield res
                break
            except DBAPIError as dpe:
                if not (retry_on_disconnect and isinstance(dpe.orig, DatabaseError)):
                    raise
                logger.warning(
                    "Trying to re-execute due to database error: %s", str(dpe)
                )
                retry_on_disconnect = False
                if session:
                    session.close()
    finally:
        if session:
            session.close()


def fetchproxy_cursor(
    sql, params=None, chunk_size=50000, sf_config=None, retry_on_disconnect=True
):
    """Execute SQL and yield results in chunks using fetchmany."""
    if params:
        sql, params = validator.format_identifiers(sql, params)

    while True:
        session = _get_sessionmaker(sf_config)()
        try:
            res = session.execute(text(sql), params)
            while True:
                chunk = res.fetchmany(chunk_size)
                if not chunk:
                    break
                yield chunk
            break
        except DBAPIError as dpe:
            if not (retry_on_disconnect and isinstance(dpe.orig, DatabaseError)):
                raise
            logger.warning("Trying to re-execute due to database error: %s", str(dpe))
            retry_on_disconnect = False
        finally:
            session.close()


def fetchproxy_chunked(
    sql, params=None, chunk_size=50000, sf_config=None, retry_on_disconnect=True
):
    """Execute SQL using LIMIT/OFFSET pagination.

    WARNING: Query must be ordered for deterministic results.
    """
    if params:
        sql, params = validator.format_identifiers(sql, params)

    offset = 0
    while True:
        # Construct paginated query
        paginated_sql = f"{sql} LIMIT {chunk_size} OFFSET {offset}"

        session = _get_sessionmaker(sf_config)()
        try:
            res = session.execute(text(paginated_sql), params)
            rows = res.fetchall()
            if not rows:
                break
            yield rows

            if len(rows) < chunk_size:
                break

            offset += chunk_size
        except DBAPIError as dpe:
            if not (retry_on_disconnect and isinstance(dpe.orig, DatabaseError)):
                raise
            logger.warning("Trying to re-execute due to database error: %s", str(dpe))
            retry_on_disconnect = False
            # For chunked, we might want to retry the current chunk,
            # but simple retry logic here restarts the specific chunk query
        finally:
            session.close()