#!/bin/bash

# Script for setting up and running SIA on RunPod with retry logic
# Requires the RunPod CLI (runpodctl) to be installed and configured

# Export MSYS_NO_PATHCONV=1 for Git Bash compatibility
export MSYS_NO_PATHCONV=1

# =================== Configuration ===================

POD_NAME="sia-agent"
CONTAINER_DISK_SIZE=200     # GB
VOLUME_SIZE=50              # GB
VOLUME_PATH="/root/data"    # Mount path within container
DOCKER_IMAGE="nvidia/cuda:12.2.0-runtime-ubuntu22.04"

GPU_CONFIGS=(
  "NVIDIA RTX A6000"
  "NVIDIA RTX 6000 Ada Generation"
)

# Load environment variables from .env file if it exists
if [ -f .env ]; then
    echo "Loading environment variables from .env file"
    source .env
fi

# Function to create pod with given parameters
create_pod() {
  local gpu_type=$1
  local cloud_arg=$2
  
  echo "Attempting to create pod with:"
  echo "- GPU Type: $gpu_type"
  echo "- Cloud Option: $cloud_arg"
  
  create_cmd="runpodctl create pod \
    --name \"$POD_NAME\" \
    --imageName \"$DOCKER_IMAGE\" \
    --gpuType \"$gpu_type\" \
    --gpuCount 1 \
    --containerDiskSize $CONTAINER_DISK_SIZE \
    --volumeSize $VOLUME_SIZE \
    --volumePath \"$VOLUME_PATH\" \
    --vcpu 1 \
    --ports \"8080/http,8888/http\" \
    --env \"SIA_REPO_URL=$SIA_REPO_URL\" \
    --env \"SIA_REPO_USER=$SIA_REPO_USER\" \
    --env \"SIA_REPO_PAT=$SIA_REPO_PAT\" \
    --env \"SIA_HF_API_KEY=$SIA_HF_API_KEY\" \
    --env \"SIA_QWQ_ENABLED=1\" \
    --env \"JUPYTER_PASSWORD=1\" \
    --args \"/bin/bash -c \\\"apt-get update && \
      apt-get install -y curl && \
      curl -Lk 'https://code.visualstudio.com/sha/download?build=stable&os=cli-alpine-x64' --output vscode_cli.tar.gz && \
      tar -xf vscode_cli.tar.gz && \
      mkdir -p /usr/local/bin && \
      mv code /usr/local/bin/ && \
      rm vscode_cli.tar.gz && \
      curl -o /root/bootstrap.sh \
      -H \"Authorization: token $SIA_REPO_PAT\" \
      https://git.nielsgeens.be/llm/SIA/raw/branch/master/scripts/bootstrap.sh && \
      yes | code tunnel \
        --name runpod \
        --install-extension ms-toolsai.jupyter \
        --install-extension ms-python.python\\\"\""
  
  # Add cloud argument if specified
  if [ -n "$cloud_arg" ]; then
    create_cmd="$create_cmd $cloud_arg"
  fi
  
  # Execute the command and capture output
  output=$(eval $create_cmd 2>&1)
  
  # Return the output
  echo "$output"
}

# Function to extract pod ID from output
extract_pod_id() {
  local output=$1
  local pod_id=""
  
  # Extract the last line of output (should contain pod ID)
  local last_line=$(echo "$output" | tail -n 1)
  
  # Extract pod ID - look for the word "pod" followed by ID in quotes
  if [[ $last_line =~ pod\ \"([a-zA-Z0-9]+)\" ]]; then
    pod_id="${BASH_REMATCH[1]}"
  else
    # Try to get the last word which is often the pod ID
    pod_id=$(echo "$last_line" | awk '{print $NF}')
  fi
  
  # If the ID has a double quote at the end, remove it
  pod_id=${pod_id%\"}
  
  # Validate we have a proper pod ID (alphanumeric)
  if [[ $pod_id =~ ^[a-zA-Z0-9]+$ ]]; then
    echo "$pod_id"
    return 0
  else
    return 1
  fi
}

# Function to wait for pod to be ready
wait_for_pod() {
  local pod_id=$1
  local max_attempts=30
  local attempt=1
  
  while [ $attempt -le $max_attempts ]; do
    echo "Waiting for pod $pod_id to be ready... (Attempt $attempt/$max_attempts)"
    sleep 5
    
    # Get the output directly and check if it contains RUNNING
    output=$(runpodctl get pod "$pod_id" 2>&1)
    
    # Print the output for debugging
    echo "$output"
    
    # Check both STATUS  RUNNING (with two spaces) and STATUS RUNNING (with one space)
    if echo "$output" | grep -q "STATUS  *RUNNING" || echo "$output" | grep -q "RUNNING"; then
      echo "Pod is now running!"
      echo "https://www.runpod.io/console/pods"
      return 0
    fi
    
    attempt=$((attempt + 1))
  done
  
  echo "Pod did not reach RUNNING state within the timeout period."
  return 1
}

# Main execution logic with retry mechanism
echo "Starting pod creation with retry logic..."

# Try different combinations of GPU types and cloud options
for gpu_type in "${GPU_CONFIGS[@]}"; do
  # Try standard cloud (no specific arg)
  echo "Trying with GPU type: $gpu_type (standard cloud)"
  output=$(create_pod "$gpu_type" "")
  echo "$output"
  
  pod_id=$(extract_pod_id "$output")
  if [ -n "$pod_id" ]; then
    echo "Successfully created pod with ID: $pod_id"
    wait_for_pod "$pod_id"
    exit 0
  fi
  
  # Try community cloud
  echo "Trying with GPU type: $gpu_type (community cloud)"
  output=$(create_pod "$gpu_type" "--communityCloud")
  echo "$output"
  
  pod_id=$(extract_pod_id "$output")
  if [ -n "$pod_id" ]; then
    echo "Successfully created pod with ID: $pod_id"
    wait_for_pod "$pod_id"
    exit 0
  fi
  
  # Try secure cloud
  echo "Trying with GPU type: $gpu_type (secure cloud)"
  output=$(create_pod "$gpu_type" "--secureCloud")
  echo "$output"
  
  pod_id=$(extract_pod_id "$output")
  if [ -n "$pod_id" ]; then
    echo "Successfully created pod with ID: $pod_id"
    wait_for_pod "$pod_id"
    exit 0
  fi
done

echo "Failed to create pod after trying all GPU types and cloud options."
exit 1