#!/bin/bash

# 设置错误处理
set -e

# 创建Hugging Face配置目录
mkdir -p ~/.huggingface

# 检查是否存在Hugging Face token
if [ ! -z "$HUGGINGFACE_TOKEN" ]; then
    echo "Setting up Hugging Face credentials..."
    echo "default_token=$HUGGINGFACE_TOKEN" > ~/.huggingface/token
fi

# 检查Android项目
if [ -d "/workspace/android-project" ]; then
    echo "Android project detected, installing dependencies..."
    cd /workspace/android-project
    if [ -f "gradlew" ]; then
        chmod +x gradlew
        ./gradlew --no-daemon dependencies
    fi
fi

# 检查CUDA可用性
python3 -c "import torch; print('CUDA available:', torch.cuda.is_available())"

# 创建示例Python脚本
cat > /workspace/test_environment.py << EOL
from transformers import pipeline
import torch

print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("CUDA device:", torch.cuda.get_device_name(0))

# 测试Hugging Face pipeline
classifier = pipeline("sentiment-analysis")
result = classifier("Hello, world!")
print("Test pipeline result:", result)
EOL

echo "Environment setup complete!"
echo "You can test the environment by running: python3 /workspace/test_environment.py"

# 执行传入的命令或启动bash
exec "$@"