#!/bin/bash
# FoundationPose deployment script (optimized for HuggingFace)

set -e

IMAGE_NAME="gpue/foundationpose-base"
TAG="latest"
PLATFORM="linux/amd64"
HF_SPACE="gpue/foundationpose"
HF_TOKEN_FILE="../training/.env.local"

echo "==================================="
echo "FoundationPose Deployment"
echo "==================================="
echo ""

# Stage 1: Build and push base image
echo "Stage 1: Building base image"
echo "Platform: ${PLATFORM}"
echo "Image: ${IMAGE_NAME}:${TAG}"
echo ""

# Check Docker login
if [ ! -f ~/.docker/config.json ] || ! grep -q "index.docker.io" ~/.docker/config.json 2>/dev/null; then
    echo "Error: Not logged in to DockerHub"
    echo "Please run: docker login"
    exit 1
fi
echo "✓ DockerHub authentication verified"
echo ""

echo "Building base image..."
docker build --platform ${PLATFORM} -f Dockerfile.base -t ${IMAGE_NAME}:${TAG} .

echo ""
echo "✓ Base image built successfully"
echo ""

# Show image size
IMAGE_SIZE=$(docker images ${IMAGE_NAME}:${TAG} --format "{{.Size}}")
echo "Image size: ${IMAGE_SIZE}"
echo ""

echo "Pushing to DockerHub..."
docker push ${IMAGE_NAME}:${TAG}

echo ""
echo "✓ Base image pushed to DockerHub: ${IMAGE_NAME}:${TAG}"
echo ""

# Stage 2: Deploy to HuggingFace
echo "Stage 2: Deploying to HuggingFace Space"
echo ""

# Initialize git repo if needed
if [ ! -d .git ]; then
    echo "Initializing git repository..."
    git init
    git remote add origin https://huggingface.co/spaces/${HF_SPACE}
    echo "✓ Git repository initialized"
    echo ""
fi

# Check if there are changes to commit
if [[ -n $(git status -s) ]]; then
    echo "Committing changes..."
    git add Dockerfile Dockerfile.base requirements.txt deploy.sh app.py client.py estimator.py
    git commit -m "Optimized Docker build to fix OOM errors"
    echo "✓ Changes committed"
else
    echo "No changes to commit"
fi

# Push to HuggingFace
echo ""
echo "Pushing to HuggingFace Space: ${HF_SPACE}"
git push https://huggingface.co/spaces/${HF_SPACE} main --force

echo ""
echo "✓ Pushed to HuggingFace"
echo ""
echo "HuggingFace will now:"
echo "  1. Pull base image from DockerHub (${IMAGE_NAME}:${TAG})"
echo "  2. Install build tools temporarily"
echo "  3. Build C++ extensions with GPU"
echo "  4. Remove build tools"
echo "  5. Download model weights (246MB)"
echo "  6. Start the Gradio app"
echo ""

# Follow build logs
echo "Following build logs..."
echo "Press Ctrl+C to stop watching"
echo ""

# Load HF token
if [ -f "${HF_TOKEN_FILE}" ]; then
    HF_TOKEN=$(grep "^HUGGINGFACE_TOKEN=" "${HF_TOKEN_FILE}" | cut -d'=' -f2)

    if [ -n "${HF_TOKEN}" ]; then
        curl -N -H "Authorization: Bearer ${HF_TOKEN}" \
            "https://huggingface.co/api/spaces/${HF_SPACE}/logs/build" 2>/dev/null | \
            while IFS= read -r line; do
                # Parse JSON and extract data field
                echo "$line" | grep -o '"data":"[^"]*"' | sed 's/"data":"//;s/"$//' | sed 's/\\n/\n/g'
            done

        echo ""
        echo "===================================="
        echo "Build Status Check"
        echo "===================================="
        echo ""

        # Wait a moment for status to update
        sleep 2

        # Check final build status
        STATUS_JSON=$(curl -s -H "Authorization: Bearer ${HF_TOKEN}" \
            "https://huggingface.co/api/spaces/${HF_SPACE}")

        STAGE=$(echo "$STATUS_JSON" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('runtime', {}).get('stage', 'UNKNOWN'))" 2>/dev/null)
        ERROR_MSG=$(echo "$STATUS_JSON" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('runtime', {}).get('errorMessage', ''))" 2>/dev/null)

        echo "Final Status: ${STAGE}"

        if [ "${STAGE}" = "RUNNING" ]; then
            echo "✓ Deployment successful!"
            echo ""
            echo "Space URL: https://${HF_SPACE/\//-}.hf.space"
            echo "API URL: https://${HF_SPACE/\//-}.hf.space/gradio_api/info"
            echo ""
            echo "Test with: cd ../training && make test-perception-api"
        elif [ "${STAGE}" = "BUILD_ERROR" ]; then
            echo "✗ Build failed!"
            if [ -n "${ERROR_MSG}" ]; then
                echo "Error: ${ERROR_MSG}"
            fi
            echo ""
            echo "If still getting OOM errors, consider:"
            echo "  - Moving weights to runtime download (not build time)"
            echo "  - Requesting larger build instance from HuggingFace"
            echo "  - Using only CUDA arch 7.5 (T4 only)"
            exit 1
        else
            echo "Status: ${STAGE}"
            if [ -n "${ERROR_MSG}" ]; then
                echo "Message: ${ERROR_MSG}"
            fi
        fi
    else
        echo "Warning: HF_TOKEN not found in ${HF_TOKEN_FILE}"
        echo "To follow logs manually:"
        echo "  curl -N -H \"Authorization: Bearer \$HF_TOKEN\" \"https://huggingface.co/api/spaces/${HF_SPACE}/logs/build\""
    fi
else
    echo "Warning: ${HF_TOKEN_FILE} not found"
    echo "To follow logs manually:"
    echo "  curl -N -H \"Authorization: Bearer \$HF_TOKEN\" \"https://huggingface.co/api/spaces/${HF_SPACE}/logs/build\""
fi