import logging

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

from dbnd import log_duration, log_metrics
from dbnd._core.errors import DatabandError
from dbnd._core.utils.timezone import utcnow
from dbnd_snowflake.snowflake_values import SnowflakeController


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",
    "TOTAL_ELAPSED_TIME": "total_elapsed_time_milliseconds",
}
RESOURCE_METRICS = ",".join(
    '"{}"'.format(m) for m in SNOWFLAKE_METRIC_TO_UI_NAME.keys()
)
SNOWFLAKE_RESULT_LIMIT_MAX_VALUE = 10000
RESULT_LIMIT_INC = 10


def log_snowflake_resource_usage(
    database: str,
    user: str,
    connection_string: str,
    query_id: str,
    session_id: Optional[int] = None,
    key: str = "snowflake_query",
    history_window: float = 15,
    query_history_result_limit: int = 100,
    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 user: Name of the user who issues the query.
    :param connection_string: Snowflake connection string to use.
    :param query_id: Supply `query_id` generated by Snowflake for more reliable search in QUERY_HISTORY.
        Recommended to use over `query_text` as search by `query_id` is more reliable.
    :param session_id: Supply `session_id` generated by Snowflake for more reliable 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 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.
    """

    if not query_id:
        error_msg = "query_id cannot be empty"
        if raise_on_error:
            raise DatabandError(error_msg)
        else:
            logger.error(error_msg)

    result_limit = min(query_history_result_limit, SNOWFLAKE_RESULT_LIMIT_MAX_VALUE)
    tries, sf_query = 0, ""
    try:
        # XXX: Do we actually need log_duration?
        with log_duration("log_snowflake_resource_usage__time_seconds", "system"):
            while result_limit <= SNOWFLAKE_RESULT_LIMIT_MAX_VALUE and tries <= retries:
                metrics_logged, sf_query = _log_snowflake_resource_usage(
                    database,
                    user,
                    connection_string,
                    query_id=query_id,
                    session_id=session_id,
                    key=key,
                    history_window=history_window,
                    query_history_result_limit=result_limit,
                )
                if metrics_logged:
                    return
                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, SNOWFLAKE_RESULT_LIMIT_MAX_VALUE
                )
                logger.info(
                    "Extending QUERY_HISTORY() search window: RESULT_LIMIT={}".format(
                        result_limit
                    )
                )
                tries += 1
                if retry_pause:
                    logger.info("Sleeping for %s", retry_pause)
                    sleep(retry_pause)

    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', user - '%s', database - '%s', connection - '%s', "
            "query - '%s'",
            query_id,
            session_id,
            user,
            database,
            conn_without_pass,
            sf_query,
        )
        if raise_on_error:
            raise

    logger.error(
        "Resource metrics were not found for query_id '%s'", query_id,
    )
    log_metrics(
        {
            key + ".warning": "No resources info found",
            # converting to str, since can be too large for DB int
            key + ".session_id": str(session_id),
            key + ".query_id": query_id,
        },
        source="user",
    )


def _build_snowflake_resource_usage_query(
    database: str,
    user: str,
    query_id: str,
    session_id: int = None,
    history_window: float = 15,
    query_history_result_limit: int = 100,
) -> str:
    time_end = utcnow()
    time_start = time_end - timedelta(minutes=history_window)
    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 LOWER(user_name)=LOWER('{user_name}') and 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,
            user_name=user,
        )
        return query_history

    if query_id:
        query_history = dedent(
            """\
            select {metrics}
            from table({database}.information_schema.query_history_by_user(
                USER_NAME => '{user_name}',
                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,
            user_name=user,
        )
        return query_history


def _log_snowflake_resource_usage(
    database,  # type: str
    user,  # type: str
    connection_string,  # type: str
    query_id,  # type: str
    session_id=None,  # type: Optional[int]
    key=None,  # type: Optional[str]
    history_window=15,  # type: float
    query_history_result_limit=100,  # type: int
):  # type: (...) -> Tuple[int, str]
    # Quick and dirty way to handle optional clause element.
    # Might be better to use SQLAlchemy expression language here
    key = key or "snowflake_query"
    query_history = _build_snowflake_resource_usage_query(
        database,
        user,
        session_id=session_id,
        query_id=query_id,
        history_window=history_window,
        query_history_result_limit=query_history_result_limit,
    )

    result = _connect_and_query(connection_string, query_history)
    if not result:
        return 0, 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
    log_metrics(metrics_to_log, source="user")
    return len(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
