import logging

from typing import List, Optional, Union

from snowflake.connector import SnowflakeConnection

from dbnd import log_duration, log_metrics
from dbnd_snowflake.snowflake_config import SnowflakeConfig
from dbnd_snowflake.snowflake_controller import SnowflakeController, SnowflakeError


logger = logging.getLogger(__name__)


def log_snowflake_resource_usage(
    database: str,
    connection_string: Union[str, SnowflakeConnection],
    query_ids: List[str],
    session_id: Optional[int] = None,
    key: str = "snowflake_query",
    history_window: float = 15,
    query_history_result_limit: Optional[int] = None,
    delay: int = 1,
    retries: int = 3,
    retry_pause: float = 0,
    raise_on_error: bool = False,
) -> None:
    """
    Search for a query previously executed by Snowflake in it's QUERY_HISTORY and log cpu time,
    run time, disk read, and other resources.

    Query's metadata can appear in QUERY_HISTORY with a lag up to 45 minutes.

    :param database: Name of the database query was issued to.
    :param connection_string: Snowflake connection string to use or existing SnowflakeConnection object.
    :param query_ids: Supply a list of `query_id` generated by Snowflake for search in QUERY_HISTORY.
    :param session_id: Supply `session_id` generated by Snowflake for more efficient search in QUERY_HISTORY.
    :param key: Override it if you call this function twice or more within the same task/Airflow Operator
    :param history_window: How deep to search into QUERY_HISTORY. Set in minutes
    :param query_history_result_limit: Passed through directly to QUERY_HISTORY search function as `RESULT_LIMIT` param
    :param delay: Initial delay before looking in QUERY_HISTORY.
        Metadata can appear there with some delay. Use this param for fine tuning
    :param retries: How much times to search in QUERY_HISTORY.
        Each time search is widened by increasing `RESULT_LIMIT` param.
    :param raise_on_error: By default all exceptions are muted so your task success status
        is not affected by errors in tracking. Set to true to re-raise all exceptions.
    :param retry_pause: Set number of seconds to pause before next retry.
    """

    with SnowflakeController(connection_string) as snowflake_ctrl:
        snowflake_config = SnowflakeConfig()
        if query_history_result_limit is None:
            query_history_result_limit = snowflake_config.query_history_result_limit

        if not all(query_id for query_id in query_ids):
            error_msg = f"query_ids cannot be empty. You supplied: {query_ids}"
            if raise_on_error:
                raise SnowflakeError(error_msg)
            else:
                logger.error(error_msg)
                return

        metrics_to_log = {}
        # XXX: Do we actually need log_duration?
        with log_duration("log_snowflake_resource_usage__time_seconds", "system"):
            for i, query_id in enumerate(query_ids):
                query_key = f"{key}.{i}" if len(query_ids) > 1 else key
                metrics_to_log.update(
                    snowflake_ctrl.get_resource_usage(
                        database,
                        query_id,
                        session_id,
                        query_key,
                        history_window,
                        query_history_result_limit,
                        delay,
                        retries,
                        retry_pause,
                        raise_on_error,
                        snowflake_config,
                    )
                )

        log_metrics(metrics_to_log, source="user")
