import logging

from datetime import timedelta
from decimal import Decimal
from textwrap import dedent
from time import sleep
from typing import Dict, List, Optional, Tuple

from dbnd import log_duration, log_metrics
from dbnd._core.utils.timezone import utcnow
from dbnd_snowflake.snowflake_config import SnowflakeConfig
from dbnd_snowflake.snowflake_values import SnowflakeController, SnowflakeError


logger = logging.getLogger(__name__)

SNOWFLAKE_METRIC_TO_UI_NAME = {
    "BYTES_SCANNED": "bytes_scanned",
    "COMPILATION_TIME": "compilation_time_milliseconds",
    "CREDITS_USED_CLOUD_SERVICES": "credits_used_cloud_services",
    "ERROR_MESSAGE": "error_message",
    "EXECUTION_STATUS": "execution_status",
    "EXECUTION_TIME": "execution_time_milliseconds",
    "QUERY_ID": "query_id",
    "QUERY_TAG": "query_tag",
    "QUERY_TEXT": "query_text",
    "ROWS_PRODUCED": "rows_produced",
    "SESSION_ID": "session_id",
    "TOTAL_ELAPSED_TIME": "total_elapsed_time_milliseconds",
}
RESOURCE_METRICS = ",".join(
    '"{}"'.format(m) for m in SNOWFLAKE_METRIC_TO_UI_NAME.keys()
)
RESULT_LIMIT_INC = 10


def log_snowflake_resource_usage(
    database: str,
    connection_string: str,
    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 = 0,
    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.
    :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.
    """

    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(
                _get_snowflake_resource_usage(
                    database,
                    connection_string,
                    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")


def _get_snowflake_resource_usage(
    database: str,
    connection_string: str,
    query_id: str,
    session_id: Optional[int],
    key: str,
    history_window: float,
    query_history_result_limit: int,
    delay: int,
    retries: int,
    retry_pause: float,
    raise_on_error: bool,
    config: SnowflakeConfig,
) -> Dict:
    if delay > 0:
        logger.info("Delaying search in QUERY_HISTORY for %s seconds", delay)
        sleep(delay)
    result_limit = min(
        query_history_result_limit, config.query_history_result_limit_max_value
    )
    tries, sf_query = 0, ""
    try:
        while (
            result_limit <= config.query_history_result_limit_max_value
            and tries <= retries
        ):
            resource_metrics, sf_query = _query_snowflake_resource_usage(
                database,
                connection_string,
                query_id=query_id,
                session_id=session_id,
                key=key,
                history_window=history_window,
                query_history_result_limit=result_limit,
                config=config,
            )
            if resource_metrics:
                return resource_metrics
            logger.warning(
                "Metadata not found for session_id '{}', query_id '{}'\n"
                "Query used to search for resource usage: '{}'".format(
                    session_id, query_id, sf_query
                )
            )
            result_limit = min(
                result_limit * RESULT_LIMIT_INC,
                config.query_history_result_limit_max_value,
            )
            logger.info(
                "Extending QUERY_HISTORY() search window: RESULT_LIMIT={}".format(
                    result_limit
                )
            )
            tries += 1
            if retry_pause and retry_pause > 0:
                logger.info("Sleeping for %s seconds", retry_pause)
                sleep(retry_pause)
        else:
            logger.info(
                "No more retries left to fetch Snoflake query resources. Giving up."
            )

    except Exception as exc:
        conn_without_pass = _censor_password(connection_string)
        logger.exception(
            "Failed to log_snowflake_resource_usage (query_text=%s, connection_string=%s)\n"
            "Last query params used to search for resource usage: query_id - '%s', "
            "sesion_id = '%s', database - '%s', connection - '%s', query - '%s'",
            query_id,
            session_id,
            database,
            conn_without_pass,
            sf_query,
        )
        if raise_on_error:
            raise

    logger.error(
        "Resource metrics were not found for query_id '%s'.\n Query used: %s",
        query_id,
        sf_query,
    )
    return {
        f"{key}.warning": "No resources info found",
        # converting to str, since can be too large for DB int
        f"{key}.session_id": str(session_id),
        f"{key}.query_id": query_id,
    }


def _build_snowflake_resource_usage_query(
    database: str,
    query_id: str,
    session_id: int,
    history_window: float,
    query_history_result_limit: int,
    config: SnowflakeConfig,
) -> str:

    time_end = utcnow() - timedelta(minutes=config.query_history_end_time_range_end)
    time_start = time_end - timedelta(
        minutes=history_window or config.query_history_end_time_range_start
    )
    if session_id:
        query_history = dedent(
            """\
            select {metrics}
            from table({database}.information_schema.query_history_by_session(
                SESSION_ID => {session_id},
                END_TIME_RANGE_START => '{time_start}'::timestamp_ltz,
                END_TIME_RANGE_END => '{time_end}'::timestamp_ltz,
                RESULT_LIMIT => {result_limit}
            ))
            where query_id='{query_id}'
            order by start_time desc limit 1;"""
        ).format(
            metrics=RESOURCE_METRICS,
            database=database,
            minutes=history_window,
            session_id=session_id,
            result_limit=query_history_result_limit,
            time_start=time_start,
            time_end=time_end,
            query_id=query_id,
        )
        return query_history

    else:
        query_history = dedent(
            """\
            select {metrics}
            from table({database}.information_schema.query_history(
                END_TIME_RANGE_START => '{time_start}'::timestamp_ltz,
                END_TIME_RANGE_END => '{time_end}'::timestamp_ltz,
                RESULT_LIMIT => {result_limit}
            ))
            where query_id='{query_id}'
            order by start_time desc limit 1;"""
        ).format(
            metrics=RESOURCE_METRICS,
            database=database,
            minutes=history_window,
            result_limit=query_history_result_limit,
            time_start=time_start,
            time_end=time_end,
            query_id=query_id,
        )
        return query_history


def _query_snowflake_resource_usage(
    database,  # type: str
    connection_string,  # type: str
    query_id,  # type: str
    session_id,  # type: Optional[int]
    key,  # type: Optional[str]
    history_window,  # type: float
    query_history_result_limit,  # type: int
    config,  # type: SnowflakeConfig
):  # type: (...) -> Tuple[Dict, str]
    key = key or "snowflake_query"
    query_history = _build_snowflake_resource_usage_query(
        database,
        query_id=query_id,
        session_id=session_id,
        history_window=history_window,
        query_history_result_limit=query_history_result_limit,
        config=config,
    )

    result = _connect_and_query(connection_string, query_history)
    if not result:
        return {}, query_history

    metrics = result[0]

    metrics_to_log = {}
    for metric, ui_name in SNOWFLAKE_METRIC_TO_UI_NAME.items():
        if metric in metrics:
            value = metrics[metric]
            # Quick hack to track decimal values. probably should be handled on a serialization level
            if isinstance(value, Decimal):
                value = float(value)
            metrics_to_log[key + "." + ui_name] = value
    return metrics_to_log, query_history


def _connect_and_query(connection_string, query, *params):
    """ connect if needed, then query. """
    if not connection_string:
        logger.error("Connection string is empty, don't know where to connect")
        return

    with SnowflakeController(connection_string) as snowflake:
        return snowflake._query(query, params)


def _censor_password(connection_string):
    """
    example connection string:
        snowflake://user:password@account.europe-west4.gcp/database
    returns:
        snowflake://user:*****@account.europe-west4.gcp/database
    """
    if (not connection_string) or ("@" not in connection_string):
        return connection_string

    split1 = connection_string.split("@")
    split2 = split1[0].split(":")

    if len(split2) != 3:
        return connection_string

    split2[-1] = "*****"
    split2_join = ":".join(split2)
    split1[0] = split2_join
    split1_join = "@".join(split1)
    return split1_join
