"""
Timeout wrapper for SAM2 initialization to prevent hanging
"""

import signal
import functools
from typing import Any, Callable


class TimeoutError(Exception):
    pass


def timeout(seconds: int = 120):
    """Decorator to add timeout to function calls"""
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            # Define signal handler
            def timeout_handler(signum, frame):
                raise TimeoutError(f"Function {func.__name__} timed out after {seconds} seconds")
            
            # Set signal handler
            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(seconds)
            
            try:
                result = func(*args, **kwargs)
            finally:
                # Restore old handler
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)
                
            return result
        return wrapper
    return decorator


@timeout(120)  # 2 minute timeout
def safe_init_state(predictor, video_path: str, **kwargs) -> Any:
    """Safely initialize SAM2 state with timeout"""
    return predictor.init_state(
        video_path=video_path,
        **kwargs
    )