# Copyright (c) 2021 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

import time

from click import echo, style

from czc.logging import get_logger

logger = get_logger()


def ok(msg):
    """Decorate f with msg, print OK if no exception, print FAILED and raise otherwise
    """
    def d(f):
        def g(*args, **kwargs):
            echo(f'{msg}: ', nl=False)
            try:
                result = f(*args, **kwargs)
                echo(style('OK', fg='green'))
                return result
            except Exception as e:
                echo(style('FAILED', fg='red'))
                logger.debug(e)
                raise (e)
        return g
    return d


def wait_for(msg, seconds=5):
    @ok(f'Waiting for {msg}')
    def f():
        time.sleep(seconds)
    f()
