#!/usr/bin/env python3
"""
Timeout Utilities for MIRA Operations
====================================

This module provides timeout mechanisms to prevent hangs in MIRA operations.
It includes decorators and context managers for adding timeouts to functions
and operations that might hang indefinitely.

Purpose:
    - Prevent startup hangs from infinite loops or blocking operations
    - Add configurable timeouts to file operations, database queries, and API calls
    - Provide graceful failure handling when operations take too long
    - Support both synchronous and asynchronous timeout patterns

Key Features:
    - Function decorator for automatic timeout handling
    - Context manager for timeout blocks
    - Configurable timeout values
    - Graceful error handling and logging
    - Support for both thread-based and signal-based timeouts

Usage:
    ```python
    from utils.timeout_utils import with_timeout, timeout_context
    
    # Decorator usage
    @with_timeout(30)  # 30 second timeout
    def slow_operation():
        # Operation that might hang
        pass
        
    # Context manager usage
    with timeout_context(15):  # 15 second timeout
        # Operation that might hang
        pass
    ```

Author: MIRA Memory System
Version: 1.0 (Timeout Prevention)
"""

import signal
import threading
import time
import functools
import logging
from typing import Any, Callable, Optional, TypeVar, Union
from contextlib import contextmanager

logger = logging.getLogger(__name__)

T = TypeVar('T')

class TimeoutError(Exception):
    """Raised when an operation times out."""
    pass


def with_timeout(timeout_seconds: float, default_return: Any = None, raise_on_timeout: bool = True):
    """
    Decorator to add timeout to a function.
    
    Args:
        timeout_seconds: Maximum time to allow function to run
        default_return: Value to return if timeout occurs and raise_on_timeout is False
        raise_on_timeout: Whether to raise TimeoutError or return default_return
        
    Returns:
        Decorated function with timeout behavior
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> T:
            result = [None]  # Use list to make it mutable in nested function
            exception = [None]
            
            def target():
                try:
                    result[0] = func(*args, **kwargs)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=target, daemon=True)
            thread.start()
            thread.join(timeout=timeout_seconds)
            
            if thread.is_alive():
                # Thread is still running, timeout occurred
                logger.warning(f"Function {func.__name__} timed out after {timeout_seconds} seconds")
                if raise_on_timeout:
                    raise TimeoutError(f"Function {func.__name__} timed out after {timeout_seconds} seconds")
                else:
                    return default_return
            
            if exception[0]:
                raise exception[0]
                
            return result[0]
        
        return wrapper
    return decorator


@contextmanager
def timeout_context(timeout_seconds: float, raise_on_timeout: bool = True):
    """
    Context manager to add timeout to a block of code.
    
    Args:
        timeout_seconds: Maximum time to allow the block to run
        raise_on_timeout: Whether to raise TimeoutError on timeout
        
    Yields:
        None
        
    Raises:
        TimeoutError: If the block times out and raise_on_timeout is True
    """
    result = [None]
    exception = [None]
    completed = [False]
    
    def target():
        try:
            yield
            completed[0] = True
        except Exception as e:
            exception[0] = e
    
    thread = threading.Thread(target=target, daemon=True)
    start_time = time.time()
    thread.start()
    thread.join(timeout=timeout_seconds)
    
    if not completed[0] and thread.is_alive():
        # Thread is still running, timeout occurred
        elapsed = time.time() - start_time
        logger.warning(f"Context block timed out after {elapsed:.1f} seconds")
        if raise_on_timeout:
            raise TimeoutError(f"Context block timed out after {elapsed:.1f} seconds")
    
    if exception[0]:
        raise exception[0]


def with_subprocess_timeout(timeout_seconds: float = 30):
    """
    Decorator specifically for subprocess operations.
    
    Args:
        timeout_seconds: Maximum time to allow subprocess to run
        
    Returns:
        Decorated function with subprocess timeout behavior
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> T:
            # Add timeout to subprocess calls if not already specified
            if 'timeout' not in kwargs:
                kwargs['timeout'] = timeout_seconds
            
            try:
                return func(*args, **kwargs)
            except Exception as e:
                if 'timeout' in str(e).lower():
                    logger.warning(f"Subprocess {func.__name__} timed out after {timeout_seconds} seconds")
                    raise TimeoutError(f"Subprocess {func.__name__} timed out after {timeout_seconds} seconds")
                raise
        
        return wrapper
    return decorator


class FileOperationTimeout:
    """Context manager for file operations with timeout."""
    
    def __init__(self, timeout_seconds: float = 10):
        self.timeout_seconds = timeout_seconds
        self.start_time = None
    
    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.start_time:
            elapsed = time.time() - self.start_time
            if elapsed > self.timeout_seconds:
                logger.warning(f"File operation took {elapsed:.1f} seconds (timeout: {self.timeout_seconds})")
    
    def check_timeout(self):
        """Check if timeout has been exceeded and raise TimeoutError if so."""
        if self.start_time:
            elapsed = time.time() - self.start_time
            if elapsed > self.timeout_seconds:
                raise TimeoutError(f"File operation timed out after {elapsed:.1f} seconds")


def safe_call_with_timeout(func: Callable, timeout_seconds: float = 30, 
                          default_return: Any = None, *args, **kwargs) -> Any:
    """
    Safely call a function with timeout, returning default on timeout or error.
    
    Args:
        func: Function to call
        timeout_seconds: Maximum time to allow function to run
        default_return: Value to return on timeout or error
        *args, **kwargs: Arguments to pass to function
        
    Returns:
        Function result or default_return on timeout/error
    """
    try:
        decorated_func = with_timeout(timeout_seconds, default_return, raise_on_timeout=False)(func)
        return decorated_func(*args, **kwargs)
    except Exception as e:
        logger.warning(f"Function {func.__name__} failed: {e}")
        return default_return


# Pre-configured timeout decorators for common operations
database_timeout = functools.partial(with_timeout, 15)  # 15 seconds for database operations
file_timeout = functools.partial(with_timeout, 10)      # 10 seconds for file operations
network_timeout = functools.partial(with_timeout, 30)   # 30 seconds for network operations
memory_timeout = functools.partial(with_timeout, 60)    # 60 seconds for memory operations