"""Script utilities."""

import functools
import time
from datetime import timedelta

import config
from logger import logger


def single_column(f):
    """Return single column query value..

    """
    def wrapper(*args, **kwargs) -> list:
        result = f(*args, **kwargs)
        return [r[0] for r in result]
    return wrapper


def handle_errors(custom_handler=None):
    """Retry on db errors.

    Args:
        custom_handler: Custom error handling function.

    """
    def inner(f):
        def wrapped(*args, **kwargs):
            for i in range(1, config.RETRY_COUNT + 1):
                try:
                    return f(*args, **kwargs)
                except Exception as ex:
                    if i != config.RETRY_COUNT:
                        logger.info(f"Retry on error: {ex}")
                        if custom_handler:
                            custom_handler(*args, **kwargs)
                        time.sleep(i * config.RETRY_WAIT_RATE)
                    else:
                        logger.info(ex)
                        raise ex

        return wrapped
    return inner


def timing(f):
    """Log execution time.

    """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = f(*args, **kwargs)
        elapsed_time = time.time() - start_time
        logger.debug(f"{f.__name__} {timedelta(seconds=elapsed_time)}")
        return result
    return wrapper