from metaflow import user_step_decorator
from metaflow.exception import MetaflowException
from metaflow.plugins.timeout_decorator import TimeoutException
import signal


@user_step_decorator
def timeout(step_name, flow, inputs=None, attributes=None):
    secs = (
        int(attributes.get("hours", 0)) * 3600
        + int(attributes.get("minutes", 0)) * 60
        + int(attributes.get("seconds", 0))
    )
    if not secs:
        raise MetaflowException("Specify a duration for @timeout.")

    def _sigalrm_handler(signum, frame):
        raise TimeoutException(f"Step {step_name} timed out after {secs} seconds.")

    signal.signal(signal.SIGALRM, _sigalrm_handler)
    signal.alarm(secs)
    yield
    signal.alarm(0)