#!/bin/bash

# Script for setting up and running SIA on RunPod
# 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 ===================
# These can be overridden by environment variables or .env file

# Pod configuration
GPU_TYPE=${GPU_TYPE:-"NVIDIA RTX 6000 Ada Generation"}
GPU_COUNT=${GPU_COUNT:-1}
CONTAINER_DISK_SIZE=${CONTAINER_DISK_SIZE:-100}     # GB
CPU_COUNT=${CPU_COUNT:-1}                           # vCPUs
POD_NAME=${POD_NAME:-"sia-agent"}
VOLUME_SIZE=${VOLUME_SIZE:-200}                     # GB
VOLUME_PATH=${VOLUME_PATH:-"/root/data"}            # Mount path within container

# Docker configuration
DOCKER_IMAGE=${DOCKER_IMAGE:-"runpod/vscode-server:0.0.0"}

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

# Create pod
echo "Creating RunPod instance with GPU type: $GPU_TYPE"
output=$(runpodctl create pod \
    --name "$POD_NAME" \
    --imageName "$DOCKER_IMAGE" \
    --gpuType "$GPU_TYPE" \
    --gpuCount $GPU_COUNT \
    --containerDiskSize $CONTAINER_DISK_SIZE \
    --volumeSize $VOLUME_SIZE \
    --volumePath "$VOLUME_PATH" \
    --vcpu $CPU_COUNT \
    --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" \
)

echo "$output"

# Extract the last line of output (should contain pod ID)
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 "Invalid pod ID detected: $pod_id" >&2
    exit 1
fi


while [ true ]; do
    echo "Waiting for pod $pod_id to be ready..."
    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!"
        break
    fi
done