#!/bin/bash
# =============================================================================
# RUNPOD SELF-TERMINATING ENTRYPOINT WRAPPER
# =============================================================================
# Purpose: Wrap entrypoint-generic.sh and automatically terminate pod on success
# Architecture:
#   1. Execute entrypoint-generic.sh (handles volume validation, binary execution)
#   2. Capture training exit code
#   3. If training succeeded (exit 0): Terminate pod via runpodctl
#   4. If training failed (exit != 0): Preserve pod for debugging
#
# Usage:
#   docker run -v runpod-volume:/runpod-volume \
#     -e RUNPOD_POD_ID=abc123 \
#     jgrusewski/foxhunt:latest \
#     /runpod-volume/binaries/train_tft_parquet \
#     --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet \
#     --epochs 50
#
# Environment Variables Required:
#   - RUNPOD_POD_ID: Auto-injected by RunPod (pod identifier)
#   - Optional: RPOD_API_KEY (if using REST API fallback)
# =============================================================================

set -euo pipefail

# Logging helper
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] WRAPPER: $*"
}

log "================================================================"
log "Foxhunt Self-Terminating Wrapper Started"
log "================================================================"
log "Pod ID: ${RUNPOD_POD_ID:-NOT_SET}"
log "Termination Strategy: runpodctl (pod-scoped)"
log "================================================================"

# Validate RUNPOD_POD_ID exists
if [ -z "${RUNPOD_POD_ID:-}" ]; then
    log "ERROR: RUNPOD_POD_ID environment variable not set"
    log "This variable is automatically injected by RunPod"
    log "If you see this error, your pod configuration may be incorrect"
    exit 1
fi

log "✓ Environment validated (RUNPOD_POD_ID present)"

# Execute the generic entrypoint script with all arguments
# This script handles volume validation, binary listing, and training execution
log "Calling entrypoint-generic.sh with args: $*"
log "================================================================"

# Run entrypoint-generic.sh and capture its exit code
# Note: We use || true to prevent set -e from exiting immediately on failure
/entrypoint-generic.sh "$@"
TRAINING_EXIT_CODE=$?

log "================================================================"
log "Training Process Completed"
log "Exit Code: $TRAINING_EXIT_CODE"
log "================================================================"

# Check if training succeeded
if [ "$TRAINING_EXIT_CODE" -eq 0 ]; then
    log "✓ TRAINING SUCCEEDED (exit code 0)"
    log "Initiating pod self-termination to save costs..."
    log ""

    # Attempt to terminate the pod using runpodctl
    # This is pre-installed in all RunPod containers with pod-scoped API key
    log "Executing: runpodctl remove pod $RUNPOD_POD_ID"

    if runpodctl remove pod "$RUNPOD_POD_ID" 2>&1 | tee /tmp/termination.log; then
        log "✓ Pod termination initiated successfully"
        log "Pod will shut down shortly..."
        log "================================================================"
        exit 0
    else
        TERM_EXIT_CODE=$?
        log "⚠ WARNING: Pod termination command failed (exit code $TERM_EXIT_CODE)"
        log "This is non-critical - training succeeded"
        log "Termination output:"
        cat /tmp/termination.log || true
        log ""
        log "Possible reasons:"
        log "  - Pod may have already been terminated"
        log "  - Network connectivity issue"
        log "  - RunPod API temporary unavailability"
        log ""
        log "Action: Pod will remain running for manual cleanup"
        log "================================================================"
        exit 0  # Training succeeded, don't fail the job
    fi
else
    log "✗ TRAINING FAILED (exit code $TRAINING_EXIT_CODE)"
    log "Preserving pod for debugging and log inspection"
    log ""
    log "Action Required:"
    log "  1. Connect to pod via SSH or web terminal"
    log "  2. Review logs in /workspace/ or container output"
    log "  3. Debug training script issues"
    log "  4. Manually terminate pod when debugging complete"
    log ""
    log "To manually terminate this pod:"
    log "  runpodctl remove pod $RUNPOD_POD_ID"
    log "================================================================"
    exit "$TRAINING_EXIT_CODE"  # Propagate failure exit code
fi