# Copyright (c) 2023 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from functools import wraps
from typing import Callable, Optional, Union

from aws_lambda_powertools import Logger

from .snowflake import connect
from .constants import PROJECT_NAME
from .user import get_user
from .session import get_session
from .warehouse import get_cluster, get_default_warehouse


logger = Logger(service=PROJECT_NAME)


def with_snowflake_connection(
    database: Union[str, Callable[[], str]], *,
    readonly: bool = True,
    warehouse: Optional[Union[str, Callable[[], str]]] = None,
) -> Callable[[Callable], Callable]:
    """Create a Snowflake connection for the given `database` and `warehouse`
    and pass it to the decorated function via the `connection` keyword argument.

    Defaults to the environment-appropriate `GENERAL` warehouse and to the `READONLY` user/role.
    """

    def snowflake_connection_handler(func: Callable) -> Callable:

        @wraps(func)
        def func_wrapper(*args, **kwargs):
            wh = warehouse() if callable(warehouse) else warehouse
            wh = wh or get_default_warehouse()
            db_name = database() if callable(database) else database

            logger.info('Connecting to Snowflake', extra={'db_name': db_name, 'readonly': readonly, 'warehouse': wh})
            conn = connect(get_cluster(),
                           db_user=get_user(readonly=readonly),
                           db_name=db_name,
                           warehouse=wh,
                           aws_session=get_session())
            logger.info('Connected to Snowflake', extra={'db_name': db_name, 'readonly': readonly, 'warehouse': wh})

            return func(*args, connection=conn, **kwargs)

        return func_wrapper

    return snowflake_connection_handler
