setup

#!/bin/bash

# BrightSign Python CV Extension - Setup Script
# Automates Steps 0-1: Prerequisites, clone, and source download

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Default Configuration
BRIGHTSIGN_OS_MAJOR_VERSION=${BRIGHTSIGN_OS_MAJOR_VERSION:-9.1}
BRIGHTSIGN_OS_MINOR_VERSION=${BRIGHTSIGN_OS_MINOR_VERSION:-52}

usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Setup BrightSign Python CV Extension development environment"
    echo ""
    echo "Options:"
    echo "  -y, --yes              Skip confirmation prompt"
    echo "  -v, --version VERSION  Set BrightSign OS version (e.g., 9.1.52)"
    echo "  --major VERSION        Set major.minor version (e.g., 9.1)"
    echo "  --minor VERSION        Set minor version number (e.g., 52)"
    echo "  -h, --help             Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0                     # Interactive setup with default version"
    echo "  $0 -y                  # Automatic setup, no prompts"
    echo "  $0 -v 9.1.53           # Setup with specific OS version"
    echo "  $0 --major 9.2 --minor 1  # Setup with version 9.2.1"
    echo ""
    echo "Default BrightSign OS version: ${BRIGHTSIGN_OS_MAJOR_VERSION}.${BRIGHTSIGN_OS_MINOR_VERSION}"
}

# Parse command line arguments
AUTO_YES=false
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -y|--yes) AUTO_YES=true; shift ;;
        -v|--version) 
            if [[ $2 =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
                # Format: major.minor.patch
                BRIGHTSIGN_OS_MAJOR_VERSION="${2%.*}"
                BRIGHTSIGN_OS_MINOR_VERSION="${2##*.}"
            elif [[ $2 =~ ^[0-9]+\.[0-9]+$ ]]; then
                # Format: major.minor
                BRIGHTSIGN_OS_MAJOR_VERSION="$2"
                BRIGHTSIGN_OS_MINOR_VERSION="0"
            else
                echo "Invalid version format: $2. Use major.minor or major.minor.patch"
                exit 1
            fi
            shift 2 ;;
        --major) 
            BRIGHTSIGN_OS_MAJOR_VERSION="$2"; shift 2 ;;
        --minor) 
            BRIGHTSIGN_OS_MINOR_VERSION="$2"; shift 2 ;;
        -h|--help)
            usage
            exit 0
            ;;
        *) echo "Unknown parameter: $1"; usage; exit 1 ;;
    esac
done

# Set the full version string
BRIGHTSIGN_OS_VERSION=${BRIGHTSIGN_OS_MAJOR_VERSION}.${BRIGHTSIGN_OS_MINOR_VERSION}

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

log() {
    echo -e "${BLUE}[$(date +'%T')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

error() {
    echo -e "${RED}❌ $1${NC}"
    exit 1
}

success() {
    echo -e "${GREEN}✅ $1${NC}"
}

# Prerequisites check
check_prerequisites() {
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log "🔍 Checking prerequisites..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Check architecture
    echo "Checking system architecture..."
    if [[ "$(uname -m)" != "x86_64" ]]; then
        error "This build requires x86_64 architecture. Current: $(uname -m)

💡 Solution: Use an x86_64 system or cloud VM (AWS, Google Cloud, Azure)
   Apple Silicon (M1/M2/M3) is NOT compatible with RKNN toolkit"
    fi
    success "Architecture: x86_64 ✓"
    echo ""

    # Check Docker
    echo "Checking Docker installation..."
    if ! command -v docker &> /dev/null; then
        error "Docker is required but not installed.

💡 Install Docker:
   • Linux: https://docs.docker.com/engine/install/
   • Mac: https://docs.docker.com/docker-for-mac/install/
   • Windows: https://docs.docker.com/docker-for-windows/install/"
    fi
    success "Docker: installed ✓"
    echo ""

    # Check Docker daemon
    echo "Checking Docker daemon..."
    if ! docker info &> /dev/null; then
        error "Docker daemon is not running.

💡 Start Docker:
   • Linux: sudo systemctl start docker
   • Mac/Windows: Start Docker Desktop application"
    fi
    success "Docker daemon: running ✓"
    echo ""

    # Check required tools
    echo "Checking required tools..."
    for tool in wget tar git; do
        if ! command -v $tool &> /dev/null; then
            error "$tool is required but not installed.

💡 Install $tool:
   • Debian/Ubuntu: sudo apt-get install $tool
   • Fedora/RHEL: sudo dnf install $tool
   • macOS: brew install $tool"
        fi
    done
    success "Required tools (wget, tar, git): installed ✓"
    echo ""

    # Check disk space (need ~25GB)
    echo "Checking disk space..."
    available_space=$(df . | awk 'NR==2 {print $4}')
    required_space=$((25 * 1024 * 1024)) # 25GB in KB

    if [[ $available_space -lt $required_space ]]; then
        warn "Warning: Less than 25GB free space available. Build may fail."
        warn "Available: $(($available_space / 1024 / 1024))GB, Recommended: 25GB+"
        echo ""
        echo "💡 Free up space by:"
        echo "   • Removing unused Docker images: docker system prune -a"
        echo "   • Clearing package caches"
        echo "   • Moving project to larger disk"
    else
        success "Disk space: $(($available_space / 1024 / 1024))GB available ✓"
    fi
    echo ""

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "✅ All prerequisites satisfied"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

# Setup Docker build environment
setup_docker() {
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log "🐳 Setting up Docker build environment..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Download Dockerfile if not present
    if [[ ! -f "Dockerfile" ]]; then
        log "📥 Downloading Dockerfile..."
        wget "https://raw.githubusercontent.com/brightsign/extension-template/refs/heads/main/Dockerfile" || \
            error "Failed to download Dockerfile

💡 Troubleshooting:
   • Check internet connection
   • Try downloading manually from:
     https://github.com/brightsign/extension-template/blob/main/Dockerfile"
    fi
    success "Dockerfile ready ✓"
    echo ""

    # Build Docker image
    log "🔨 Building Docker image..."
    echo "⏱️  This will take 5-10 minutes (downloads BrightSign OS source)"
    echo "💡 Progress: Watch for 'Step X/Y' messages below"
    echo ""

    docker build --rm \
        --build-arg USER_ID=$(id -u) \
        --build-arg GROUP_ID=$(id -g) \
        --build-arg BRIGHTSIGN_OS_VERSION=${BRIGHTSIGN_OS_VERSION} \
        --ulimit memlock=-1:-1 \
        -t bsoe-build . || error "Failed to build Docker image

💡 Common issues:
   • Docker daemon not running → Start Docker service
   • Insufficient disk space → Free up at least 25GB
   • Network timeout → Retry with better internet connection
   • Permission denied → Ensure user is in 'docker' group:
     sudo usermod -aG docker \$USER && newgrp docker"

    success "Docker image built ✓"
    echo ""

    # Create srv directory
    log "📁 Creating workspace directories..."
    mkdir -p srv
    success "Workspace directories created ✓"
    echo ""

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "✅ Docker environment ready"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

# Clone rknn-toolkit2 repository
clone_rknn_toolkit() {
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log "📦 Cloning rknn-toolkit2 repository..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Create toolkit directory if it doesn't exist
    mkdir -p toolkit

    # Clone or update rknn-toolkit2
    if [[ -d "toolkit/rknn-toolkit2" ]]; then
        log "📂 rknn-toolkit2 already exists, updating..."
        cd toolkit/rknn-toolkit2
        git fetch || warn "⚠️  Failed to fetch updates (using existing version)"
        git checkout v2.3.2 || warn "⚠️  Failed to checkout v2.3.2 tag (using current version)"
        cd ../..
    else
        log "📥 Cloning rknn-toolkit2 from GitHub..."
        echo "⏱️  This may take 2-3 minutes depending on your connection"
        echo ""
        cd toolkit
        git clone https://github.com/airockchip/rknn-toolkit2.git || error "Failed to clone rknn-toolkit2

💡 Troubleshooting:
   • Check internet connection
   • Verify GitHub is accessible
   • Try cloning manually:
     cd toolkit
     git clone https://github.com/airockchip/rknn-toolkit2.git"
        cd rknn-toolkit2
        git checkout v2.3.2 || error "Failed to checkout v2.3.2 tag"
        cd ../..
        success "rknn-toolkit2 cloned successfully ✓"
    fi
    echo ""

    # Verify the wheel file exists (v2.3.2)
    log "🔍 Verifying wheel file..."
    WHEEL_PATH="toolkit/rknn-toolkit2/rknn-toolkit-lite2/packages/rknn_toolkit_lite2-2.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl"
    if [[ -f "$WHEEL_PATH" ]]; then
        success "rknn-toolkit-lite2 wheel file found ✓"
    else
        error "Expected wheel file not found at: $WHEEL_PATH

💡 This usually means the repository was not cloned correctly.
   Try removing and re-cloning:
   rm -rf toolkit/rknn-toolkit2
   ./setup"
    fi
    echo ""

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "✅ rknn-toolkit2 ready"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

# Clone rknn_model_zoo repository
clone_rknn_model_zoo() {
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log "📦 Cloning rknn_model_zoo repository..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Clone or update rknn_model_zoo
    if [[ -d "toolkit/rknn_model_zoo" ]]; then
        log "📂 rknn_model_zoo already exists, updating..."
        cd toolkit/rknn_model_zoo
        git fetch || warn "⚠️  Failed to fetch updates (using existing version)"
        git checkout v2.3.2 || warn "⚠️  Failed to checkout v2.3.2 tag (using current version)"
        cd ../..
    else
        log "📥 Cloning rknn_model_zoo from GitHub..."
        echo "⏱️  This may take 2-3 minutes depending on your connection"
        echo ""
        cd toolkit
        git clone https://github.com/airockchip/rknn_model_zoo.git || error "Failed to clone rknn_model_zoo

💡 Troubleshooting:
   • Check internet connection
   • Verify GitHub is accessible
   • Try cloning manually:
     cd toolkit
     git clone https://github.com/airockchip/rknn_model_zoo.git"
        cd rknn_model_zoo
        git checkout v2.3.2 || error "Failed to checkout v2.3.2 tag"
        cd ../..
        success "rknn_model_zoo cloned successfully ✓"
    fi
    echo ""

    # Verify yolox example exists
    log "🔍 Verifying model zoo examples..."
    YOLOX_PATH="toolkit/rknn_model_zoo/examples/yolox/python"
    if [[ -d "$YOLOX_PATH" ]]; then
        success "rknn_model_zoo examples found ✓"
    else
        error "Expected yolox example not found at: $YOLOX_PATH

💡 This usually means the repository was not cloned correctly.
   Try removing and re-cloning:
   rm -rf toolkit/rknn_model_zoo
   ./setup"
    fi
    echo ""

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "✅ rknn_model_zoo ready (50+ NPU-optimized models available)"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

# Main setup function
main() {
    echo "BrightSign Python CV Extension - Setup Script"
    echo "============================================="
    echo "This script will:"
    echo "1. Check prerequisites"
    echo "2. Setup Docker build environment"
    echo "3. Clone rknn-toolkit2 repository"
    echo "4. Clone rknn_model_zoo repository"
    echo "5. Download BrightSign OS source (~2-3 GB) inside Docker container"
    echo ""
    echo "Building for BrightSign OS version: ${BRIGHTSIGN_OS_VERSION}"
    echo ""
    
    # Confirm execution if AUTO_YES is not set
    if [[ "$AUTO_YES" != true ]]; then
        read -p "Continue? [y/N] " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            echo "Setup cancelled."
            exit 0
        fi
    else
        log "Automatic mode: skipping confirmation prompt"
    fi
    
    start_time=$(date +%s)
    
    check_prerequisites
    setup_docker
    clone_rknn_toolkit
    clone_rknn_model_zoo
    
    end_time=$(date +%s)
    duration=$((end_time - start_time))

    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "🎉 SETUP COMPLETE"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "Completed in $(($duration / 60))m $(($duration % 60))s"
    echo ""
    echo "📋 Next Steps:"
    echo ""
    echo "1. Build the SDK (30-60 minutes):"
    echo "   ./build --extract-sdk"
    echo ""
    echo "2. Package the extension (5 minutes):"
    echo "   ./package"
    echo ""
    echo "3. Deploy to player:"
    echo "   See docs/deployment.md for complete instructions"
    echo ""
    echo "4. Test with Model Zoo:"
    echo "   See docs/model-zoo-guide.md for NPU inference examples"
    echo ""
    echo "📚 Documentation:"
    echo "   • Quick start: QUICKSTART.md"
    echo "   • Command reference: WORKFLOWS.md"
    echo "   • FAQ: FAQ.md"
    echo "   • Getting started guide: docs/getting-started.md"
    echo ""
    echo "💡 Environment variables:"
    echo "   export BRIGHTSIGN_OS_VERSION=${BRIGHTSIGN_OS_VERSION}"
    echo "   export project_root=$(pwd)"
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi