#!/usr/bin/env python3
"""
Cross-platform test wrapper with memory and timeout limits.
Prevents runaway tests from consuming all system resources.

MEMORY LIMITING STRATEGY:
- Linux: Uses resource.RLIMIT_AS/RLIMIT_RSS (works reliably)
- macOS: Active monitoring with process killing (resource limits often fail)
- Windows: Active monitoring only (no resource module)

Usage:
    python3 run_with_memory_limit.py <memory_mb> <timeout_sec> <test_exe> [args...]

Exit Codes:
    0   - Test passed
    1-123 - Test failed (exit code from test)
    124 - Timeout exceeded
    137 - Out of memory (killed by monitor)

Examples:
    python3 run_with_memory_limit.py 4608 180 ./my_test
    python3 run_with_memory_limit.py 2048 60 ./large_test --verbose
"""

import sys
import os
import subprocess
import signal
import time
import threading

def set_memory_limit(memory_mb):
    """
    Set memory limit for the current process and all child processes.

    Args:
        memory_mb: Memory limit in megabytes

    Returns:
        True if limit was set successfully, False otherwise

    Note:
        - On Linux: Uses resource.RLIMIT_AS (virtual memory) unless ASAN is detected
        - On macOS: Uses resource.RLIMIT_DATA (data segment)
        - On Windows: No-op (returns False, but continues execution)
        - With ASAN: Skips virtual memory limits (ASAN needs huge shadow memory reservation)
    """
    try:
        import resource
        import platform
        memory_bytes = memory_mb * 1024 * 1024

        # Check if AddressSanitizer is enabled
        # ASAN needs huge virtual address space for shadow memory (~30TB on ARM64)
        # but doesn't actually use that much physical RAM
        asan_enabled = os.environ.get('ASAN_OPTIONS') is not None

        if asan_enabled:
            # Don't set virtual memory limits with ASAN - rely on RSS monitoring only
            print(f"Note: ASAN detected - skipping virtual memory limits, using RSS monitoring only",
                  file=sys.stderr)
            return False

        # macOS has issues with RLIMIT_AS, so use RLIMIT_DATA instead
        # RLIMIT_DATA limits the data segment size (heap memory)
        if platform.system() == "Darwin":  # macOS
            limit_type = resource.RLIMIT_DATA
            limit_name = "RLIMIT_DATA"
        else:  # Linux and others
            limit_type = resource.RLIMIT_AS
            limit_name = "RLIMIT_AS"

        # Get current limits to check if our limit is valid
        soft, hard = resource.getrlimit(limit_type)

        # If hard limit is set and lower than our desired limit, use hard limit
        if hard != resource.RLIM_INFINITY and hard < memory_bytes:
            print(f"Warning: Requested memory limit ({memory_mb}MB) exceeds system limit ({hard // (1024*1024)}MB)",
                  file=sys.stderr)
            print(f"Using system limit instead", file=sys.stderr)
            memory_bytes = hard

        # Set the limit
        resource.setrlimit(limit_type, (memory_bytes, memory_bytes))

        return True

    except (ImportError, AttributeError) as e:
        # Windows doesn't have resource module
        print(f"Warning: Memory limiting not available on this platform: {e}",
              file=sys.stderr)
        return False
    except (OSError, ValueError) as e:
        # Permission issues or invalid limit
        # This is not fatal - test can still run with timeout protection
        print(f"Warning: Could not set memory limit: {e}", file=sys.stderr)
        return False

def get_process_memory_mb(pid):
    """
    Get process memory usage in MB (works on macOS, Linux, and with psutil).

    Returns:
        Memory usage in MB, or None if unable to determine
    """
    try:
        # Try psutil first (most reliable if installed)
        import psutil
        process = psutil.Process(pid)
        # RSS = Resident Set Size (actual physical memory used)
        return process.memory_info().rss / (1024 * 1024)
    except ImportError:
        pass
    except Exception:
        return None

    # Fallback: Parse ps output (works on Unix/macOS)
    try:
        # ps -o rss= returns RSS in KB
        result = subprocess.run(
            ['ps', '-o', 'rss=', '-p', str(pid)],
            capture_output=True,
            text=True,
            timeout=1
        )
        if result.returncode == 0:
            rss_kb = int(result.stdout.strip())
            return rss_kb / 1024.0
    except Exception:
        pass

    return None

def monitor_memory(proc, memory_limit_mb, check_interval=0.5):
    """
    Monitor process memory usage and kill if limit exceeded.
    Runs in background thread.

    Args:
        proc: subprocess.Popen object
        memory_limit_mb: Memory limit in MB
        check_interval: How often to check memory (seconds)

    Returns:
        True if process was killed for exceeding memory, False otherwise
    """
    killed_for_memory = False

    while proc.poll() is None:  # While process is still running
        try:
            mem_mb = get_process_memory_mb(proc.pid)

            if mem_mb is not None and mem_mb > memory_limit_mb:
                print(f"\nERROR: Test exceeded memory limit!", file=sys.stderr)
                print(f"  Limit: {memory_limit_mb} MB", file=sys.stderr)
                print(f"  Used:  {mem_mb:.1f} MB", file=sys.stderr)
                print(f"  Killing process {proc.pid}...", file=sys.stderr)

                proc.kill()  # SIGKILL - immediate termination
                killed_for_memory = True
                break

            time.sleep(check_interval)

        except Exception as e:
            # Process may have exited between poll() and memory check
            break

    return killed_for_memory

def run_with_monitoring(cmd, timeout_sec, memory_limit_mb):
    """
    Run command with active memory monitoring and timeout protection.

    This is used when resource limits cannot be set (e.g., macOS, Windows).

    Args:
        cmd: Command and arguments as list
        timeout_sec: Maximum execution time in seconds
        memory_limit_mb: Maximum memory in MB

    Returns:
        Exit code from command, or 124/137 for timeout/OOM
    """
    try:
        # Start process
        proc = subprocess.Popen(
            cmd,
            stdout=sys.stdout,
            stderr=sys.stderr
        )

        # Start memory monitor thread
        monitor_result = [False]  # Use list to allow modification in thread

        def monitor_thread():
            monitor_result[0] = monitor_memory(proc, memory_limit_mb)

        monitor_t = threading.Thread(target=monitor_thread, daemon=True)
        monitor_t.start()

        # Wait for process with timeout
        try:
            returncode = proc.wait(timeout=timeout_sec)
            monitor_t.join(timeout=1.0)  # Wait briefly for monitor to finish

            if monitor_result[0]:
                return 137  # Killed for memory
            return returncode

        except subprocess.TimeoutExpired:
            print(f"\nERROR: Test exceeded {timeout_sec} second timeout",
                  file=sys.stderr)
            proc.kill()
            proc.wait()
            return 124

    except KeyboardInterrupt:
        print(f"\nTest interrupted by user", file=sys.stderr)
        if proc.poll() is None:
            proc.kill()
            proc.wait()
        return 130

    except Exception as e:
        print(f"\nERROR: Unexpected error running test: {e}", file=sys.stderr)
        if proc.poll() is None:
            proc.kill()
            proc.wait()
        return 1

def main():
    """Main entry point"""
    if len(sys.argv) < 4:
        print("Usage: run_with_memory_limit.py <memory_mb> <timeout_sec> <test_exe> [args...]",
              file=sys.stderr)
        print("\nExample:", file=sys.stderr)
        print("  python3 run_with_memory_limit.py 4608 180 ./my_test", file=sys.stderr)
        sys.exit(1)

    try:
        memory_mb = int(sys.argv[1])
        timeout_sec = int(sys.argv[2])
        test_cmd = sys.argv[3:]

        if memory_mb <= 0 or timeout_sec <= 0:
            print("ERROR: Memory and timeout must be positive integers", file=sys.stderr)
            sys.exit(1)

        if not test_cmd:
            print("ERROR: No test command specified", file=sys.stderr)
            sys.exit(1)

    except ValueError as e:
        print(f"ERROR: Invalid arguments: {e}", file=sys.stderr)
        print("Memory and timeout must be integers", file=sys.stderr)
        sys.exit(1)

    # Try to set resource limits (works on Linux, usually fails on macOS)
    limits_set = set_memory_limit(memory_mb)

    if not limits_set:
        # Resource limits failed - use active monitoring instead
        # This is the common case on macOS
        print(f"Using active memory monitoring (checking every 0.5s)", file=sys.stderr)
        exit_code = run_with_monitoring(test_cmd, timeout_sec, memory_mb)
    else:
        # Resource limits worked - still monitor as backup
        exit_code = run_with_monitoring(test_cmd, timeout_sec, memory_mb)

    sys.exit(exit_code)

if __name__ == "__main__":
    main()
