#!/bin/bash
# deploy.sh — One-command setup for vast.ai GPU instances
#
# TWO ways to use this:
#
# Option A — Download from your private HuggingFace repo + run:
#   export HF_TOKEN=your_token
#   pip install huggingface_hub
#   python -c "from huggingface_hub import snapshot_download; snapshot_download('YOUR_USER/td-toolkit', local_dir='.')"
#   bash deploy.sh demo_autopilot.td
#
# Option B — Already uploaded files manually:
#   bash deploy.sh my_pipeline.td

set -e  # Stop on any error

# Colors for pretty output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

echo ""
echo "==========================================="
echo "  TD Deploy — vast.ai GPU Setup"
echo "==========================================="
echo ""

# Check if a .td file was provided
if [ -z "$1" ]; then
    echo -e "${RED}ERROR: No .td file specified${NC}"
    echo ""
    echo "Usage: bash deploy.sh my_pipeline.td"
    echo ""
    echo "Available .td files:"
    ls -1 *.td td_lang/examples/*.td 2>/dev/null || echo "  (none found)"
    exit 1
fi

TD_FILE="$1"

if [ ! -f "$TD_FILE" ]; then
    echo -e "${RED}ERROR: File not found: $TD_FILE${NC}"
    exit 1
fi

echo -e "${GREEN}[1/5]${NC} Installing td_lang dependencies..."
pip install lark --quiet 2>/dev/null || pip install lark
echo "  Done."

# Check for HF token
echo ""
echo -e "${GREEN}[2/5]${NC} Checking environment..."
if [ -z "$HF_TOKEN" ]; then
    echo -e "${YELLOW}  WARNING: HF_TOKEN not set.${NC}"
    echo "  Models won't download from HuggingFace without it."
    echo "  Set it with: export HF_TOKEN=your_token_here"
    echo ""
    read -p "  Continue anyway? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
else
    echo "  HF_TOKEN: set"
fi

# Check td_lang is accessible
echo ""
echo -e "${GREEN}[3/5]${NC} Checking td_lang..."
if python -c "import td_lang" 2>/dev/null; then
    VERSION=$(python -c "import td_lang; print(td_lang.__version__)" 2>/dev/null || echo "unknown")
    echo "  td_lang v$VERSION: found"
else
    # Try adding current directory to path
    export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$(pwd)"
    if python -c "import td_lang" 2>/dev/null; then
        VERSION=$(python -c "import td_lang; print(td_lang.__version__)" 2>/dev/null || echo "unknown")
        echo "  td_lang v$VERSION: found (added to PYTHONPATH)"
    else
        echo -e "${RED}  ERROR: td_lang not found!${NC}"
        echo "  Make sure the td_lang/ folder is in the current directory."
        echo "  Current directory: $(pwd)"
        echo "  Contents:"
        ls -1
        exit 1
    fi
fi

# Check for rclone (needed for save command)
echo ""
echo -e "${GREEN}[4/5]${NC} Checking tools..."
if command -v rclone &> /dev/null; then
    echo "  rclone: installed"
    if rclone listremotes 2>/dev/null | grep -q "gdrive:"; then
        echo "  Google Drive: configured"
    else
        echo -e "${YELLOW}  Google Drive: not configured${NC}"
        echo "  Run 'rclone config' to set up Google Drive (name it 'gdrive')"
    fi
else
    echo -e "${YELLOW}  rclone: not installed (installing...)${NC}"
    curl -s https://rclone.org/install.sh | bash 2>/dev/null || {
        echo -e "${YELLOW}  Could not install rclone. 'save' commands won't work.${NC}"
    }
fi

# Check GPU
if command -v nvidia-smi &> /dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)
    echo "  GPU: $GPU_NAME ($GPU_MEM)"
else
    echo -e "${YELLOW}  WARNING: No GPU detected (nvidia-smi not found)${NC}"
fi

# Run the .td file
echo ""
echo -e "${GREEN}[5/5]${NC} Running: $TD_FILE"
echo "==========================================="
echo ""

python -m td_lang run "$TD_FILE"

echo ""
echo "==========================================="
echo -e "${GREEN}  TD Deploy complete!${NC}"
echo "==========================================="