#!/bin/bash
# ROBUST ERROR HANDLING - Continue on non-critical errors
set -uo pipefail  # Undefined variables are errors, pipefail catches pipe failures
# Note: NOT using 'set -e' to avoid exiting on non-critical warnings

# =============================================================================
# CRASH LOG CAPTURE FOR RUNPOD DEBUGGING
# =============================================================================
# Captures all errors and outputs them to stdout BEFORE container restarts
# This ensures crash logs are visible in RunPod console even during restart loops

# Crash log file (writable location)
CRASH_LOG="/tmp/foxhunt-crash.log"
touch "$CRASH_LOG" 2>/dev/null || CRASH_LOG="/dev/null"

# Timestamp function for all log messages
timestamp() {
    date '+%Y-%m-%d %H:%M:%S'
}

# Error handler - captures line number and command
handle_error() {
    local exit_code=$1
    local line_number=$2
    local timestamp=$(timestamp)

    echo "[$timestamp] ERROR: Command failed with exit code $exit_code at line $line_number" | tee -a "$CRASH_LOG"
    echo "[$timestamp] Last command context:" | tee -a "$CRASH_LOG"
    echo "[$timestamp]   Script: $0" | tee -a "$CRASH_LOG"
    echo "[$timestamp]   PWD: $(pwd)" | tee -a "$CRASH_LOG"
}

# Exit handler - outputs full crash log before container dies
handle_exit() {
    local exit_code=$?
    local timestamp=$(timestamp)

    if [ $exit_code -ne 0 ]; then
        echo "" | tee -a "$CRASH_LOG"
        echo "============================================" | tee -a "$CRASH_LOG"
        echo "[$timestamp] CRASH DETECTED - Exit Code: $exit_code" | tee -a "$CRASH_LOG"
        echo "============================================" | tee -a "$CRASH_LOG"

        if [ -f "$CRASH_LOG" ] && [ -s "$CRASH_LOG" ]; then
            echo "[$timestamp] Full crash log contents:" | tee -a "$CRASH_LOG"
            echo "--------------------------------------------"
            cat "$CRASH_LOG"
            echo "--------------------------------------------"
        else
            echo "[$timestamp] No crash log data captured"
        fi

        echo "[$timestamp] Container will restart in 5 seconds..." | tee -a "$CRASH_LOG"
        echo "[$timestamp] COPY THE ABOVE LOGS FOR DEBUGGING" | tee -a "$CRASH_LOG"
        echo "============================================"
    fi
}

# Set up trap handlers
trap 'handle_error $? $LINENO' ERR
trap 'handle_exit' EXIT

# Log entrypoint start
echo "[$(timestamp)] Entrypoint script started" | tee -a "$CRASH_LOG"

# =============================================================================
# RUNPOD ENTRYPOINT SCRIPT - VOLUME MOUNT ARCHITECTURE
# =============================================================================
# Executes training binary from mounted Runpod Network Volume
# NO DOWNLOADS - all binaries and data are pre-uploaded to the volume
#
# ARCHITECTURE:
#   - Docker image: CUDA runtime environment ONLY (~2GB)
#   - Runpod Network Volume: Mounted at /runpod-volume/ with:
#     * /runpod-volume/binaries/train_tft_parquet (and other training binaries)
#     * /runpod-volume/test_data/*.parquet (9 data files, 180-day datasets)
#   - Pod: Mounts volume directly (instant access, zero network transfers)
#
# Environment Variables:
#   BINARY_NAME              - Training binary name (default: train_tft_parquet)
#                              Options: train_tft_parquet, train_mamba2_parquet, train_dqn, train_ppo
#
# GPU Hardware:
#   - Tesla V100-PCIE-16GB (16GB VRAM, 5120 CUDA cores, 640 Tensor cores)
#   - Compute Capability: 7.0
#   - Memory Bandwidth: 900 GB/s
#   - FP32 Performance: 15.7 TFLOPS
#   - FP16 Performance: 125 TFLOPS (with Tensor cores)
#   - Cost: $0.10/hour
# =============================================================================

echo "=========================================="
echo "Foxhunt HFT - Runpod Volume Mount Training"
echo "=========================================="

# Default binary name if not specified
BINARY_NAME=${BINARY_NAME:-train_tft_parquet}

echo "Configuration:"
echo "  Binary: ${BINARY_NAME}"
echo "  Volume Path: /runpod-volume/"
echo "  Architecture: Direct volume mount (NO downloads)"

# =============================================================================
# VERIFY VOLUME IS MOUNTED
# =============================================================================

echo ""
echo "=========================================="
echo "Verifying Runpod Network Volume Mount"
echo "=========================================="

if [ ! -d "/runpod-volume" ]; then
    echo "[$(timestamp)] ERROR: /runpod-volume directory does not exist" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Runpod Network Volume is NOT mounted!" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Solution:" | tee -a "$CRASH_LOG"
    echo "  1. Go to Runpod console" | tee -a "$CRASH_LOG"
    echo "  2. Edit pod configuration" | tee -a "$CRASH_LOG"
    echo "  3. Add volume mount: /runpod-volume" | tee -a "$CRASH_LOG"
    echo "  4. Select your Runpod Network Volume from dropdown" | tee -a "$CRASH_LOG"
    echo "  5. Redeploy pod" | tee -a "$CRASH_LOG"
    exit 1
fi

if [ ! -d "/runpod-volume/binaries" ]; then
    echo "[$(timestamp)] ERROR: /runpod-volume/binaries directory does not exist" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Volume is mounted but binaries directory is missing!" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Solution:" | tee -a "$CRASH_LOG"
    echo "  1. Upload training binaries to Runpod Network Volume" | tee -a "$CRASH_LOG"
    echo "  2. Create directory: /runpod-volume/binaries/" | tee -a "$CRASH_LOG"
    echo "  3. Upload binaries:" | tee -a "$CRASH_LOG"
    echo "     - train_tft_parquet" | tee -a "$CRASH_LOG"
    echo "     - train_mamba2_parquet" | tee -a "$CRASH_LOG"
    echo "     - train_dqn" | tee -a "$CRASH_LOG"
    echo "     - train_ppo" | tee -a "$CRASH_LOG"
    exit 1
fi

if [ ! -d "/runpod-volume/test_data" ]; then
    echo "[$(timestamp)] ERROR: /runpod-volume/test_data directory does not exist" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Volume is mounted but test_data directory is missing!" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Solution:" | tee -a "$CRASH_LOG"
    echo "  1. Upload training data to Runpod Network Volume" | tee -a "$CRASH_LOG"
    echo "  2. Create directory: /runpod-volume/test_data/" | tee -a "$CRASH_LOG"
    echo "  3. Upload Parquet files (9 files):" | tee -a "$CRASH_LOG"
    echo "     - ES_FUT_180d.parquet" | tee -a "$CRASH_LOG"
    echo "     - NQ_FUT_180d.parquet" | tee -a "$CRASH_LOG"
    echo "     - 6E_FUT_180d.parquet" | tee -a "$CRASH_LOG"
    echo "     - ZN_FUT_90d.parquet" | tee -a "$CRASH_LOG"
    echo "     - (and 5 additional small test files)" | tee -a "$CRASH_LOG"
    exit 1
fi

echo "✓ Volume mounted successfully at /runpod-volume/"

# =============================================================================
# LOAD CREDENTIALS FROM .env FILE (IF EXISTS)
# =============================================================================

echo ""
echo "=========================================="
echo "Loading Credentials"
echo "=========================================="

if [ -f /runpod-volume/.env ]; then
  # Set secure permissions (600 = owner read/write only)
  # Ignore errors - permissions might already be correct
  chmod 600 /runpod-volume/.env 2>/dev/null || true

  # Load environment variables from .env file
  set -a  # Export all variables
  source /runpod-volume/.env || {
    echo "ERROR: Failed to source /runpod-volume/.env"
    exit 1
  }
  set +a

  echo "✓ Loaded credentials from /runpod-volume/.env"
  PERMS=$(stat -c%a /runpod-volume/.env 2>/dev/null || stat -f%Lp /runpod-volume/.env 2>/dev/null || echo "unknown")
  echo "  File permissions: ${PERMS} (expected: 600)"

  # Verify key environment variables loaded (without printing values)
  # Note: These are optional for ML training - not required for Runpod
  if [ -n "${DATABASE_URL:-}" ]; then
    echo "  ✓ DATABASE_URL loaded"
  else
    echo "  ⚠ WARNING: DATABASE_URL not set (optional for training)"
  fi

  if [ -n "${REDIS_URL:-}" ]; then
    echo "  ✓ REDIS_URL loaded"
  else
    echo "  ⚠ WARNING: REDIS_URL not set (optional for training)"
  fi

  if [ -n "${VAULT_ADDR:-}" ]; then
    echo "  ✓ VAULT_ADDR loaded"
  else
    echo "  ⚠ WARNING: VAULT_ADDR not set (optional for training)"
  fi
else
  echo "⚠ WARNING: .env file not found at /runpod-volume/.env"
  echo "  Some services may fail without credentials"
  echo ""
  echo "Solution:"
  echo "  1. Upload .env file to Runpod Network Volume:"
  echo "     aws s3 cp .env s3://foxhunt-ml-volume/.env \\"
  echo "       --endpoint-url \$RUNPOD_S3_ENDPOINT \\"
  echo "       --profile runpod"
  echo "  2. Redeploy pod"
  echo ""
  echo "Continuing with environment defaults..."
fi

# List volume contents for verification
echo ""
echo "Volume contents:"
echo "  Binaries:"
if [ -d "/runpod-volume/binaries" ]; then
    ls -lh /runpod-volume/binaries/ 2>/dev/null || echo "    (empty or not accessible)"
else
    echo "    (directory not found)"
fi

echo "  Data files:"
if [ -d "/runpod-volume/test_data" ]; then
    ls -lh /runpod-volume/test_data/*.parquet 2>/dev/null || echo "    (no parquet files found)"
else
    echo "    (directory not found)"
fi

# =============================================================================
# VERIFY TRAINING BINARY EXISTS
# =============================================================================

echo ""
echo "=========================================="
echo "Verifying Training Binary"
echo "=========================================="

# CRITICAL FIX: Detect if first argument is already a full binary path
# This allows users to override via command: "/runpod-volume/binaries/train_dqn --args..."
# Instead of defaulting to BINARY_NAME environment variable
if [ $# -gt 0 ] && [[ "$1" =~ ^/runpod-volume/binaries/ ]]; then
    # First argument is a full binary path - use it directly
    BINARY_PATH="$1"
    BINARY_NAME=$(basename "$BINARY_PATH")
    shift  # Remove binary path from arguments, rest are training params
    
    echo "Using binary path from command argument:"
    echo "  Binary: ${BINARY_NAME}"
    echo "  Full path: ${BINARY_PATH}"
    echo "  Training args: $*"
else
    # Use BINARY_NAME environment variable (default behavior)
    BINARY_PATH="/runpod-volume/binaries/${BINARY_NAME}"
    
    echo "Using binary from BINARY_NAME environment variable:"
    echo "  Binary: ${BINARY_NAME}"
    echo "  Full path: ${BINARY_PATH}"
    echo "  Training args: $*"
fi

echo ""
if [ ! -f "${BINARY_PATH}" ]; then
    echo "[$(timestamp)] ERROR: Training binary not found: ${BINARY_PATH}" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Available binaries in /runpod-volume/binaries/:" | tee -a "$CRASH_LOG"
    ls -lh /runpod-volume/binaries/ 2>&1 | tee -a "$CRASH_LOG" || echo "  (none)" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "Solution:" | tee -a "$CRASH_LOG"
    echo "  1. Upload ${BINARY_NAME} to Runpod Network Volume" | tee -a "$CRASH_LOG"
    echo "  2. Upload path: /runpod-volume/binaries/${BINARY_NAME}" | tee -a "$CRASH_LOG"
    echo "  3. Make executable: chmod +x /runpod-volume/binaries/${BINARY_NAME}" | tee -a "$CRASH_LOG"
    echo "  4. Redeploy pod" | tee -a "$CRASH_LOG"
    exit 1
fi

# Check if binary is executable
if [ ! -x "${BINARY_PATH}" ]; then
    echo "[$(timestamp)] WARNING: Binary is not executable, attempting to fix..." | tee -a "$CRASH_LOG"
    chmod +x "${BINARY_PATH}" 2>&1 | tee -a "$CRASH_LOG" || {
        echo "[$(timestamp)] ERROR: Failed to make binary executable" | tee -a "$CRASH_LOG"
        echo "       Manual fix required on volume" | tee -a "$CRASH_LOG"
        exit 1
    }
    echo "[$(timestamp)] ✓ Binary is now executable" | tee -a "$CRASH_LOG"
fi

BINARY_SIZE=$(stat -c%s "${BINARY_PATH}" 2>/dev/null || stat -f%z "${BINARY_PATH}")
echo "✓ Binary found: ${BINARY_PATH}"
echo "  Size: $(numfmt --to=iec-i --suffix=B ${BINARY_SIZE} 2>/dev/null || echo ${BINARY_SIZE} bytes)"

# =============================================================================
# VERIFY GPU HARDWARE (Tesla V100)
# =============================================================================

echo ""
echo "=========================================="
echo "GPU Hardware Status"
echo "=========================================="

if command -v nvidia-smi &> /dev/null; then
    echo "GPU Information:"
    # Make GPU queries non-fatal - capture errors but continue
    nvidia-smi --query-gpu=name,memory.total,memory.free,driver_version,compute_cap --format=csv,noheader 2>&1 || {
        echo "  WARNING: nvidia-smi query failed (GPU might be initializing)"
    }

    echo ""
    echo "Expected Hardware: Tesla V100-PCIE-16GB / RTX A4000"
    echo "  V100: 16GB VRAM, 5120 CUDA cores, Compute 7.0"
    echo "  A4000: 16GB VRAM, 6144 CUDA cores, Compute 8.6"

    echo ""
    echo "CUDA/cuDNN Status:"
    nvidia-smi -L 2>&1 || echo "  WARNING: nvidia-smi -L failed"
else
    echo "WARNING: nvidia-smi not found - GPU may not be accessible"
    echo "         Training may fall back to CPU (very slow)"
fi

# =============================================================================
# START TRAINING
# =============================================================================

echo ""
echo "=========================================="
echo "Starting Training..."
echo "=========================================="
echo "Binary: ${BINARY_PATH}"
echo "Arguments: $*"
echo "Architecture: Direct volume mount (zero network overhead)"
echo ""

# Verify binary is accessible before exec
if [ ! -f "${BINARY_PATH}" ]; then
    echo "[$(timestamp)] FATAL ERROR: Binary disappeared: ${BINARY_PATH}" | tee -a "$CRASH_LOG"
    ls -la /runpod-volume/binaries/ 2>&1 | tee -a "$CRASH_LOG" || echo "(binaries directory not accessible)" | tee -a "$CRASH_LOG"
    exit 1
fi

if [ ! -x "${BINARY_PATH}" ]; then
    echo "[$(timestamp)] FATAL ERROR: Binary is not executable: ${BINARY_PATH}" | tee -a "$CRASH_LOG"
    ls -la "${BINARY_PATH}" 2>&1 | tee -a "$CRASH_LOG"
    exit 1
fi

# Print exact command being executed for debugging
echo "[$(timestamp)] Executing command:" | tee -a "$CRASH_LOG"
echo "  ${BINARY_PATH} $*" | tee -a "$CRASH_LOG"
echo "" | tee -a "$CRASH_LOG"
echo "[$(timestamp)] If you see this message, entrypoint script completed successfully." | tee -a "$CRASH_LOG"
echo "[$(timestamp)] Any errors below are from the training binary itself." | tee -a "$CRASH_LOG"
echo "" | tee -a "$CRASH_LOG"

# Execute training binary with stderr/stdout capture
# Using process substitution to capture both streams AND exit code
echo "[$(timestamp)] Training binary output:" | tee -a "$CRASH_LOG"
echo "--------------------------------------------" | tee -a "$CRASH_LOG"

# Run binary with stderr/stdout redirect to both console and crash log
"${BINARY_PATH}" "$@" 2>&1 | tee -a "$CRASH_LOG"

# Capture exit code from the training binary (PIPESTATUS[0] gets exit code before tee)
EXIT_CODE=${PIPESTATUS[0]}

echo "--------------------------------------------" | tee -a "$CRASH_LOG"
echo "[$(timestamp)] Training binary exited with code: $EXIT_CODE" | tee -a "$CRASH_LOG"

# If binary crashed, output full crash log before exiting
if [ $EXIT_CODE -ne 0 ]; then
    echo "" | tee -a "$CRASH_LOG"
    echo "============================================" | tee -a "$CRASH_LOG"
    echo "[$(timestamp)] TRAINING BINARY CRASH DETECTED" | tee -a "$CRASH_LOG"
    echo "============================================" | tee -a "$CRASH_LOG"
    echo "[$(timestamp)] Binary: ${BINARY_PATH}" | tee -a "$CRASH_LOG"
    echo "[$(timestamp)] Exit Code: $EXIT_CODE" | tee -a "$CRASH_LOG"
    echo "[$(timestamp)] Arguments: $*" | tee -a "$CRASH_LOG"
    echo "" | tee -a "$CRASH_LOG"
    echo "[$(timestamp)] FULL CRASH LOG:" | tee -a "$CRASH_LOG"
    echo "--------------------------------------------"
    cat "$CRASH_LOG"
    echo "--------------------------------------------"
    echo "[$(timestamp)] END CRASH LOG" | tee -a "$CRASH_LOG"
    echo "============================================"

    # Exit with the same code as the training binary
    exit $EXIT_CODE
fi

echo "[$(timestamp)] Training completed successfully!" | tee -a "$CRASH_LOG"
exit 0