# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Example usage:
#
# Basic usage (will retry failed step 3 times):
# steps:
#   - uses: ultralytics/actions/retry@main
#     with:
#       run: python train.py
#
# Advanced usage:
# steps:
#   - uses: ultralytics/actions/retry@main
#     with:
#       run: |
#         python setup.py install
#         pytest tests/
#       retries: 2               # Retry twice after initial attempt (3 total runs)
#       timeout_minutes: 30      # Maximum time for each attempt
#       retry_delay_seconds: 10  # Base delay between retries in seconds
#       backoff: exponential     # exponential (10s, 20s, 40s, ...) or fixed
#       jitter: true             # Randomize delay to 80-120% to avoid thundering herd
#       shell: bash              # Use python or bash shell

name: "Step-Level Retry"
description: "Retries a step while preserving its full context"
inputs:
  timeout_minutes:
    description: "Maximum time in minutes for each attempt"
    required: false
    default: "360"
  retries:
    description: "Number of retry attempts after initial run"
    required: false
    default: "3"
  retry_delay_seconds:
    description: "Base delay between retries in seconds"
    required: false
    default: "10"
  backoff:
    description: "Backoff strategy: exponential (base * 2^n) or fixed"
    required: false
    default: "exponential"
  jitter:
    description: "Randomize delay to 80-120% of computed value to spread retries and avoid thundering herd"
    required: false
    default: "true"
  run:
    description: "Command to run"
    required: true
  shell:
    description: "Shell to use (bash or python)"
    required: false
    default: "bash"

runs:
  using: "composite"
  steps:
    - name: Execute with retry
      shell: bash
      env:
        # Note: Workflow env vars are automatically inherited.
        RETRY_COMMANDS: ${{ inputs.run }}
        RETRY_RUNNER_OS: ${{ runner.os }}
      run: |
        set +e  # Don't exit on error - we handle errors manually

        timeout_seconds=$(( ${{ inputs.timeout_minutes }} * 60 ))
        attempt=1
        max_attempts=$(( 1 + ${{ inputs.retries }} ))  # Initial run + retries
        retries=${{ inputs.retries }}
        retry_delay=${{ inputs.retry_delay_seconds }}
        backoff="${{ inputs.backoff }}"
        jitter="${{ inputs.jitter }}"
        timeout_mins=${{ inputs.timeout_minutes }}
        shell_type="${{ inputs.shell }}"
        exit_code=0

        # Validate backoff strategy
        if [ "$backoff" != "exponential" ] && [ "$backoff" != "fixed" ]; then
            echo "::error::Invalid backoff '$backoff'; must be 'exponential' or 'fixed'"
            exit 1
        fi

        # Validate numeric inputs are non-negative integers
        case "$retries$retry_delay$timeout_mins" in
            ''|*[!0-9]*) echo "::error::retries, retry_delay_seconds, timeout_minutes must be non-negative integers"; exit 1 ;;
        esac

        if { [ "$RETRY_RUNNER_OS" != "Windows" ] || [ "$shell_type" = "python" ]; } && ! command -v python3 >/dev/null; then
            echo "::error::Python 3 is required for timeout supervision on this runner"
            exit 1
        fi

        # Create temporary script file with appropriate extension
        if [ "$shell_type" = "python" ]; then
            if ! TEMP_SCRIPT=$(mktemp --suffix=.py 2>/dev/null); then
                # Fallback: create temp file then rename it
                if ! TEMP_SCRIPT=$(mktemp 2>/dev/null); then
                    echo "::error::Failed to create temporary file"
                    exit 1
                fi
                if ! mv "$TEMP_SCRIPT" "${TEMP_SCRIPT}.py"; then
                    rm -f "$TEMP_SCRIPT"
                    echo "::error::Failed to rename temporary file"
                    exit 1
                fi
                TEMP_SCRIPT="${TEMP_SCRIPT}.py"
            fi
        else
            if ! TEMP_SCRIPT=$(mktemp 2>/dev/null); then
                echo "::error::Failed to create temporary file"
                exit 1
            fi
        fi

        # Ensure cleanup on exit
        trap 'rm -f "$TEMP_SCRIPT" "${TEMP_SCRIPT}.timeout"' EXIT

        # Write the user's commands to the temp script
        printf '%s\n' "$RETRY_COMMANDS" > "$TEMP_SCRIPT"

        # Validate script was written
        if [ ! -s "$TEMP_SCRIPT" ]; then
            echo "::error::No commands provided to execute"
            exit 1
        fi

        while [ "$attempt" -le "$max_attempts" ]; do
            if [ "$attempt" -gt 1 ]; then
                echo "::group::Retry $((attempt - 1)) of $retries"
            fi

            # Execute the script with a deadline, terminating its process group and tracked descendants on timeout.
            timeout_marker="${TEMP_SCRIPT}.timeout"
            rm -f "$timeout_marker"
            if [ "$RETRY_RUNNER_OS" = "Windows" ]; then
                retry_python=""
                [ "$shell_type" = "python" ] && retry_python=$(cygpath -w "$(command -v python3)")
                RETRY_EXECUTABLE=$(cygpath -w "$BASH") \
                RETRY_PYTHON="$retry_python" \
                RETRY_SCRIPT=$(cygpath -w "$TEMP_SCRIPT") \
                RETRY_SHELL_TYPE="$shell_type" \
                RETRY_TIMEOUT_MARKER=$(cygpath -w "$timeout_marker") \
                RETRY_TIMEOUT_SECONDS="$timeout_seconds" \
                powershell.exe -NoProfile -Command '
        $startInfo = New-Object System.Diagnostics.ProcessStartInfo
        if ($env:RETRY_SHELL_TYPE -eq "python") {
            $startInfo.FileName = $env:RETRY_PYTHON
            $startInfo.Arguments = "`"$env:RETRY_SCRIPT`""
        } else {
            $startInfo.FileName = $env:RETRY_EXECUTABLE
            $startInfo.Arguments = "-e `"$env:RETRY_SCRIPT`""
        }
        $startInfo.UseShellExecute = $false
        $process = New-Object System.Diagnostics.Process
        $process.StartInfo = $startInfo
        [void]$process.Start()
        if (-not $process.WaitForExit([int]$env:RETRY_TIMEOUT_SECONDS * 1000)) {
            [System.IO.File]::WriteAllText($env:RETRY_TIMEOUT_MARKER, "")
            $allProcesses = @(Get-CimInstance Win32_Process)
            $descendants = @()
            $parents = @($process.Id)
            while ($parents.Count -gt 0) {
                $children = @($allProcesses | Where-Object { $parents -contains [int]$_.ParentProcessId })
                $parents = @($children | ForEach-Object { [int]$_.ProcessId })
                $descendants += $parents
            }
            foreach ($id in @($process.Id) + $descendants) {
                taskkill.exe /PID $id /T /F 2>$null | Out-Null
                Stop-Process -Id $id -Force -ErrorAction SilentlyContinue
            }
            exit 124
        }
        $process.WaitForExit()
        exit $process.ExitCode
        '
            else
                python3 - "$shell_type" "$TEMP_SCRIPT" "$timeout_seconds" "$timeout_marker" <<'PY'
        import os
        import signal
        import subprocess
        import sys
        import time
        import uuid

        def descendants(root_pid):
            """Return all current descendants of a process."""
            children = {}
            for entry in os.scandir("/proc"):
                if entry.name.isdigit():
                    try:
                        with open(f"{entry.path}/stat") as file:
                            stat = file.read()
                        children.setdefault(int(stat[stat.rfind(")") + 2 :].split()[1]), []).append(int(entry.name))
                    except (FileNotFoundError, PermissionError, ValueError):
                        pass
            found, pending = set(), [root_pid]
            while pending:
                for pid in children.get(pending.pop(), []):
                    if pid not in found:
                        found.add(pid)
                        pending.append(pid)
            return found

        def tracked_processes(token):
            """Return processes carrying the attempt's unique inherited environment marker."""
            marker = f"ULTRALYTICS_RETRY_TOKEN={token}"
            processes = subprocess.check_output(["ps", "axeww", "-o", "pid=", "-o", "command="], text=True)
            return {int(line.split(None, 1)[0]) for line in processes.splitlines() if marker in line}

        def send_signal(process_group, process_ids, sig):
            """Signal the process group and descendants that detached from it."""
            try:
                os.killpg(process_group, sig)
            except (PermissionError, ProcessLookupError):
                pass
            for pid in process_ids:
                try:
                    os.kill(pid, sig)
                except (PermissionError, ProcessLookupError):
                    pass

        shell_type, script, timeout, timeout_marker = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
        command = [sys.executable, script] if shell_type == "python" else ["bash", "-e", script]
        token = uuid.uuid4().hex
        environment = os.environ.copy()
        environment["ULTRALYTICS_RETRY_TOKEN"] = token
        if sys.platform.startswith("linux"):
            import ctypes

            if ctypes.CDLL(None, use_errno=True).prctl(36, 1, 0, 0, 0) != 0:
                raise OSError(ctypes.get_errno(), "Unable to enable child subreaper")
        process = subprocess.Popen(command, env=environment, start_new_session=True)
        process_ids = set()
        try:
            return_code = process.wait(timeout=timeout)
            sys.exit(128 - return_code if return_code < 0 else return_code)
        except subprocess.TimeoutExpired:
            open(timeout_marker, "w").close()
            if sys.platform.startswith("linux"):
                process_ids = descendants(os.getpid())
            process_ids |= tracked_processes(token)
            send_signal(process.pid, process_ids, signal.SIGTERM)
            deadline = time.monotonic() + 5
            while time.monotonic() < deadline:
                if sys.platform.startswith("linux"):
                    process_ids = descendants(os.getpid())
                else:
                    process_ids = set()
                process_ids |= tracked_processes(token)
                process_ids = {pid for pid in process_ids if pid != os.getpid()}
                if not process_ids:
                    break
                send_signal(process.pid, process_ids, signal.SIGTERM)
                time.sleep(0.1)
            if sys.platform.startswith("linux"):
                process_ids = descendants(os.getpid())
            else:
                process_ids = set()
            process_ids |= tracked_processes(token)
            send_signal(process.pid, process_ids, signal.SIGKILL)
            process.wait()
            sys.exit(124)
        PY
            fi
            exit_code=$?

            if [ -f "$timeout_marker" ]; then
                echo "::error::Attempt timed out after $timeout_mins minutes"
            fi

            if [ "$exit_code" -eq 0 ]; then
                [ "$attempt" -gt 1 ] && echo "::endgroup::"
                exit 0
            fi

            # Handle failure
            if [ "$attempt" -gt 1 ]; then
                echo "Retry $((attempt - 1)) failed with exit code $exit_code"
                echo "::endgroup::"
            else
                echo "Initial attempt failed with exit code $exit_code"
            fi

            if [ "$attempt" -eq "$max_attempts" ]; then
                echo "::error::Step failed after initial attempt and $((attempt - 1)) retries"
                exit "$exit_code"
            fi

            # Compute delay for the upcoming retry
            if [ "$backoff" = "exponential" ]; then
                delay=$((retry_delay * (2 ** (attempt - 1))))
            else
                delay=$retry_delay
            fi

            # Apply jitter: random 80-120% of computed delay (±20%)
            if [ "$jitter" = "true" ] && [ "$delay" -gt 0 ]; then
                low=$((delay * 8 / 10))
                span=$((delay * 4 / 10))
                rand=$((RANDOM % (span + 1)))
                delay=$((low + rand))
                [ "$delay" -lt 1 ] && delay=1  # integer division floor; never collapse to 0
            fi

            echo "Retrying in $delay seconds..."
            sleep "$delay"
            attempt=$((attempt + 1))
        done