#!/bin/bash
set -e

# Build and run script for NVIDIA Jetson devices (Orin Nano, Xavier, etc.)
# Uses Dockerfile.jetson with L4T PyTorch base image

show_usage() {
    echo "Usage: $0 [COMMAND]"
    echo ""
    echo "Commands:"
    echo "  build    - Build the Jetson Docker image"
    echo "  run      - Run the container with GPU support (default)"
    echo "  both     - Build and run"
    echo "  shell    - Run with bash shell for debugging"
    echo "  test     - Test GPU access in container"
    echo ""
    echo "Environment variables:"
    echo "  THEMATIC_MODEL_NAME  - Model to use (default: all-mpnet-base-v2)"
    echo "                         Use all-MiniLM-L6-v2 for lower GPU memory usage"
    echo ""
    echo "Example with smaller model:"
    echo "  THEMATIC_MODEL_NAME=all-MiniLM-L6-v2 $0 run"
    echo ""
}

IMAGE_NAME="crossword-app:jetson"

# Model options:
#   all-mpnet-base-v2 (420MB, best quality, default)
#   all-MiniLM-L6-v2 (90MB, faster, use if GPU memory issues)
MODEL_NAME="${THEMATIC_MODEL_NAME:-all-mpnet-base-v2}"

# GPU access for Jetson requires --runtime nvidia (not --gpus all)
DOCKER_ARGS="--rm -p 7860:7860 --runtime nvidia \
    -e ENABLE_DEBUG_TAB=true \
    -e VOCAB_SOURCE=norvig \
    -e DIFFICULTY_WEIGHT=0.2 \
    -e THEMATIC_MODEL_NAME=$MODEL_NAME"

build_image() {
    echo "🔨 Building Jetson Docker image..."
    docker build -f Dockerfile.jetson -t $IMAGE_NAME .
}

run_container() {
    echo "🚀 Running on Jetson with GPU..."
    docker run $DOCKER_ARGS $IMAGE_NAME
}

run_shell() {
    echo "🐚 Running shell for debugging..."
    docker run -it $DOCKER_ARGS $IMAGE_NAME /bin/bash
}

test_gpu() {
    echo "🔍 Testing GPU access in container..."
    docker run --rm --runtime nvidia $IMAGE_NAME python3 -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA device: {torch.cuda.get_device_name(0)}')
    # Test actual GPU memory allocation
    x = torch.randn(1000, 1000, device='cuda')
    print(f'GPU tensor created: {x.shape} on {x.device}')
    print('✅ GPU is working!')
else:
    print('❌ CUDA not available')
"
}

# Parse command
COMMAND="${1:-run}"

case "$COMMAND" in
    build)
        build_image
        ;;
    run)
        run_container
        ;;
    both)
        build_image
        run_container
        ;;
    shell)
        run_shell
        ;;
    test)
        test_gpu
        ;;
    -h|--help|help)
        show_usage
        exit 0
        ;;
    *)
        echo "Error: Unknown command '$COMMAND'"
        show_usage
        exit 1
        ;;
esac