package

#!/bin/bash

# BrightSign Python CV Extension - Package Creation Script
# Automates Step 3: Extension packaging and deployment preparation

set -e

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

# 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}[PACKAGE] $1${NC}"
}

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

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

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

# Configuration
TIMESTAMP=$(date +%Y%m%d-%H%M%S)
DEVELOPMENT_PACKAGE="pydev-${TIMESTAMP}.zip"
EXTENSION_PACKAGE="ext_pydev-${TIMESTAMP}.zip"

usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Package BrightSign Python CV Extension for deployment"
    echo ""
    echo "Options:"
    echo "  -d, --dev-only     Create development package only"
    echo "  -e, --ext-only     Create extension package only"
    echo "  -c, --clean        Clean install directory before packaging"
    echo "  -v, --verify       Run validation after packaging"
    echo "  --yolox            Include YOLOX example and test data"
    echo "  -h, --help         Show this help message"
    echo ""
    echo "Package types:"
    echo "  Development: For /usr/local deployment (volatile, testing)"
    echo "  Extension:   For permanent installation (production)"
}

# Parse command line arguments
DEV_ONLY=false
EXT_ONLY=false
CLEAN=false
VERIFY=false
INCLUDE_YOLOX=false

while [[ $# -gt 0 ]]; do
    case $1 in
        -d|--dev-only)
            DEV_ONLY=true
            shift
            ;;
        -e|--ext-only)
            EXT_ONLY=true
            shift
            ;;
        -c|--clean)
            CLEAN=true
            shift
            ;;
        -v|--verify)
            VERIFY=true
            shift
            ;;
        --yolox)
            INCLUDE_YOLOX=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Check prerequisites
check_prerequisites() {
    log "Checking prerequisites..."

    # Check if SDK exists
    if [[ ! -d "sdk" ]]; then
        error "SDK directory not found. Please extract the SDK first."
    fi

    if [[ ! -d "sdk/sysroots/aarch64-oe-linux" ]]; then
        error "SDK sysroot not found. Run: ./build --extract-sdk && ./brightsign-x86_64-cobra-toolchain-*.sh -d ./sdk -y"
    fi

    # Check for required scripts
    if [[ ! -f "sh/make-extension-lvm" ]]; then
        error "make-extension-lvm script not found"
    fi

    # Check for user init examples
    if [[ ! -d "user-init" ]]; then
        warn "user-init directory not found"
    fi

    success "Prerequisites check passed"
}

# Create install directory structure
create_install_structure() {
    log "Creating extension package structure..."

    if [[ "$CLEAN" == "true" ]]; then
        log "Cleaning existing install directory..."
        rm -rf install
    fi

    # Create directory structure
    mkdir -p install/{usr/{bin,lib},sh}

    success "Directory structure created"
}

# Copy SDK components
copy_sdk_components() {
    local sdk_sysroot="sdk/sysroots/aarch64-oe-linux"


    log "Copying Python runtime and libraries from SDK..."

    # Copy Python binaries
    log "Copying Python binaries..."
    # Use cp -P to preserve symlinks and handle them properly
    cp -P "$sdk_sysroot/usr/bin/python3" install/usr/bin/ 2>/dev/null || true
    cp -P "$sdk_sysroot/usr/bin/python3.8" install/usr/bin/ 2>/dev/null || true
    cp -P "$sdk_sysroot/usr/bin/python3.8-config-lib" install/usr/bin/ 2>/dev/null || true
    cp -P "$sdk_sysroot/usr/bin/pip3" install/usr/bin/ 2>/dev/null || true
    cp -P "$sdk_sysroot/usr/bin/pip3.8" install/usr/bin/ 2>/dev/null || true

    # Fix symlinks to be relative instead of absolute
    if [[ -L "install/usr/bin/python3" ]]; then
        cd install/usr/bin
        ln -sf python3.8 python3
        ln -sf python3.8-config-lib python3.8-config
        ln -sf python3.8-config python3-config
        cd - > /dev/null
    fi

    # Ensure librknnrt.so is present (fallback if not in SDK)
    if [ ! -f "$sdk_sysroot/usr/lib/librknnrt.so" ]; then
        warn "librknnrt.so not found in SDK, downloading..."
        if wget -q -O "$sdk_sysroot/usr/lib/librknnrt.so" \
            https://github.com/airockchip/rknn-toolkit2/raw/v2.3.2/rknpu2/runtime/Linux/librknn_api/aarch64/librknnrt.so; then
            success "librknnrt.so downloaded successfully"
        else
            error "Failed to download librknnrt.so - NPU functionality may not work"
        fi
    else
        log "librknnrt.so found in SDK"
    fi

    # Copy libraries (including librknnrt.so)
    log "Copying libraries (this may take a few minutes)..."
    cp -r "$sdk_sysroot/usr/lib" install/usr/

    # Copy setup files (test scripts now in user-init/)
    log "Copying support files..."
    # Note: setup_python_env is copied with other scripts in add_extension_scripts()

    success "SDK components copied"
}

# Add extension scripts
add_extension_scripts() {
    log "Adding extension management scripts..."

    # Copy and make executable
    cp sh/bsext_init install/ && chmod +x install/bsext_init
    cp sh/uninstall.sh install/ && chmod +x install/uninstall.sh

    # Copy supporting scripts
    mkdir -p install/sh
    cp sh/{init-extension,cleanup-extension,check-status,run-user-init,setup_python_env,pydev-env,test_python_imports} install/sh/
    chmod +x install/sh/*

    success "Extension scripts added"
}

# Copy rknn-toolkit-lite2 wheel file
copy_rknn_wheel() {
    log "Installing RKNN toolkit into extension site-packages..."

    # Install rknn-toolkit-lite2 (lightweight runtime for on-device inference)
    # Note: Full rknn-toolkit2 is NOT compatible with BrightSign due to hardcoded
    # /usr/lib64/ paths. Model zoo examples work via patched rknn_executor.py
    local lite_wheel_path="toolkit/rknn-toolkit2/rknn-toolkit-lite2/packages/rknn_toolkit_lite2-2.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl"

    local site_packages="install/usr/lib/python3.8/site-packages"
    mkdir -p "$site_packages"

    if [[ ! -f "$lite_wheel_path" ]]; then
        error "rknn-toolkit-lite2 wheel not found at: $lite_wheel_path"
        error "Run ./setup first to clone rknn-toolkit2"
        return 1
    fi

    log "Installing rknn-toolkit-lite2..."
    local temp_dir=$(mktemp -d)

    # Extract wheel contents (wheel is just a ZIP file)
    unzip -q "$lite_wheel_path" -d "$temp_dir" || {
        error "Failed to extract rknn-toolkit-lite2 wheel"
        rm -rf "$temp_dir"
        return 1
    }

    # Copy all package directories
    for pkg_dir in "$temp_dir"/*/ ; do
        local dir_name=$(basename "$pkg_dir")
        # Skip .dist-info and other metadata directories for now
        if [[ ! "$dir_name" =~ \.dist-info$ ]] && [[ ! "$dir_name" =~ ^__ ]]; then
            if [[ -d "$pkg_dir" ]]; then
                cp -r "$pkg_dir" "$site_packages/"
                log "  Installed package: $dir_name"
            fi
        fi
    done

    # Copy metadata
    for dist_info in "$temp_dir"/*.dist-info; do
        if [[ -d "$dist_info" ]]; then
            cp -r "$dist_info" "$site_packages/"
            log "  Installed metadata: $(basename "$dist_info")"
        fi
    done

    # Cleanup temporary directory
    rm -rf "$temp_dir"

    success "rknn-toolkit-lite2 installed successfully (requires OS 9.1.79.3+)"
    log "Provides 'rknnlite.api.RKNNLite' for on-device NPU inference"
}

# Copy user-init examples (including patched py_utils)
copy_user_init_examples() {
    log "Copying user-init examples..."

    local examples_src="user-init/examples"
    local examples_dst="install/examples"

    if [[ ! -d "$examples_src" ]]; then
        warn "user-init/examples directory not found - skipping"
        return 0
    fi

    mkdir -p "$examples_dst"

    # Copy all example files and directories
    cp -r "$examples_src"/* "$examples_dst/"

    # Verify py_utils was copied
    if [[ -d "$examples_dst/py_utils" ]]; then
        success "User-init examples copied (including patched py_utils for model_zoo)"
        log "py_utils provides RKNNLite compatibility wrapper for model_zoo examples"
    else
        warn "py_utils directory not found in user-init examples"
    fi
}

# Copy YOLOX example if requested
copy_yolox_example() {
    if [[ "$INCLUDE_YOLOX" == "true" ]]; then
        log "Including YOLOX example and test data..."
        
        # Check if rknn_model_zoo exists
        if [[ ! -d "toolkit/rknn_model_zoo" ]]; then
            error "rknn_model_zoo not found. Run ./setup first to clone the repository."
        fi
        
        # Create examples directory structure
        mkdir -p install/usr/lib/python3.8/examples/yolox
        
        # Copy YOLOX python example
        local yolox_src="toolkit/rknn_model_zoo/examples/yolox/python"
        if [[ -d "$yolox_src" ]]; then
            cp -r "$yolox_src"/* install/usr/lib/python3.8/examples/yolox/
            success "YOLOX Python example copied"
        else
            error "YOLOX example not found at: $yolox_src"
        fi
        
        # Copy py_utils (required dependencies)
        local pyutils_src="toolkit/rknn_model_zoo/py_utils"
        if [[ -d "$pyutils_src" ]]; then
            cp -r "$pyutils_src" install/usr/lib/python3.8/examples/
            success "py_utils dependencies copied"
        else 
            error "py_utils not found at: $pyutils_src"
        fi
        
        # Copy model download script and labels
        local model_src="toolkit/rknn_model_zoo/examples/yolox/model"
        if [[ -d "$model_src" ]]; then
            mkdir -p install/usr/lib/python3.8/examples/yolox/model
            cp "$model_src/download_model.sh" install/usr/lib/python3.8/examples/yolox/model/ 2>/dev/null || true
            cp "$model_src/coco_80_labels_list.txt" install/usr/lib/python3.8/examples/yolox/model/ 2>/dev/null || true
            success "YOLOX model files copied"
        fi
        
        # Copy a bus test image from COCO dataset examples
        local dataset_src="toolkit/rknn_model_zoo/datasets"
        if [[ -d "$dataset_src" ]]; then
            # Find and copy a bus image if available
            mkdir -p install/usr/lib/python3.8/examples/yolox/test_images
            find "$dataset_src" -name "*.jpg" -o -name "*.png" | head -3 | while read img; do
                cp "$img" install/usr/lib/python3.8/examples/yolox/test_images/ 2>/dev/null || true
            done
            success "Test images copied (if available)"
        fi
        
        # Make scripts executable
        chmod +x install/usr/lib/python3.8/examples/yolox/model/download_model.sh 2>/dev/null || true
        
        # Create a wrapper script for running YOLOX
        create_yolox_wrapper
        
        log "YOLOX example integration completed"
    fi
}

# Create YOLOX wrapper script
create_yolox_wrapper() {
    local wrapper_path="install/usr/lib/python3.8/examples/yolox/run_yolox_test.sh"
    
    cat > "$wrapper_path" << 'EOF'
#!/bin/bash

# YOLOX Test Runner for BrightSign Python CV Extension
# This script runs YOLOX inference on test images and saves results

set -e

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

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

log() {
    echo -e "${BLUE}[YOLOX] $1${NC}"
}

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

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

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

# Setup environment
setup_environment() {
    log "Setting up YOLOX test environment..."
    
    # Ensure Python environment is available
    if ! command -v python3 >/dev/null 2>&1; then
        error "Python3 not found. Ensure Python CV extension is properly installed."
    fi
    
    # Add py_utils to Python path
    export PYTHONPATH="$EXAMPLES_DIR:$PYTHONPATH"
    
    # Create output directory in user-init examples folder
    mkdir -p /usr/local/user-init/examples/yolox_results
    
    success "Environment setup complete"
}

# Download model if needed
download_model() {
    log "Checking for YOLOX model..."
    
    cd "$SCRIPT_DIR/model"
    
    if [[ ! -f "yolox_s.onnx" && ! -f "yolox_m.onnx" ]]; then
        log "Downloading YOLOX model..."
        if [[ -f "download_model.sh" ]]; then
            bash download_model.sh
        else
            warn "download_model.sh not found, you may need to manually download the model"
        fi
    else
        success "YOLOX model already available"
    fi
    
    cd - > /dev/null
}

# Convert ONNX model to RKNN
convert_model() {
    log "Converting ONNX model to RKNN format..."
    
    cd "$SCRIPT_DIR"
    
    if [[ ! -f "yolox.rknn" ]]; then
        if [[ -f "convert.py" ]]; then
            python3 convert.py
            success "Model converted to RKNN format"
        else
            warn "convert.py not found, using ONNX model for inference"
        fi
    else
        success "RKNN model already available"
    fi
    
    cd - > /dev/null
}

# Run inference on test images
run_inference() {
    log "Running YOLOX inference..."
    
    cd "$SCRIPT_DIR"
    
    # Find test images
    local test_images_dir="test_images"
    local output_dir="/usr/local/user-init/examples/yolox_results"
    
    if [[ -d "$test_images_dir" ]]; then
        for img in "$test_images_dir"/*.{jpg,png,jpeg}; do
            if [[ -f "$img" ]]; then
                local img_name=$(basename "$img")
                log "Processing: $img_name"
                
                # Run YOLOX inference
                python3 yolox.py \
                    --model_path="${PWD}/yolox.rknn" \
                    --img_path="$img" \
                    --img_show \
                    --img_save \
                    --output_dir="$output_dir" || {
                    warn "RKNN inference failed, trying ONNX..."
                    
                    # Fallback to ONNX if available
                    if [[ -f "model/yolox_s.onnx" ]]; then
                        python3 yolox.py \
                            --model_path="${PWD}/model/yolox_s.onnx" \
                            --img_path="$img" \
                            --img_show \
                            --img_save \
                            --output_dir="$output_dir"
                    fi
                }
                
                success "Processed: $img_name"
            fi
        done
    else
        warn "No test images found in $test_images_dir"
        log "You can add test images to: $SCRIPT_DIR/$test_images_dir/"
    fi
    
    cd - > /dev/null
}

# Main function
main() {
    echo "YOLOX Test Runner for BrightSign"
    echo "================================"
    
    setup_environment
    download_model
    convert_model
    run_inference
    
    echo ""
    success "YOLOX testing completed!"
    echo ""
    echo "Results saved to: /usr/local/user-init/examples/yolox_results/"
    echo "View output images and inference logs in the results directory."
}

# Run main function
main "$@"
EOF

    chmod +x "$wrapper_path"
    success "YOLOX wrapper script created"
}

# Verify installation structure
verify_structure() {
    log "Verifying installation structure..."

    local errors=0

    # Check essential directories
    for dir in "usr/bin" "usr/lib" "sh"; do
        if [[ ! -d "install/$dir" ]]; then
            error "Missing directory: install/$dir"
            ((errors++))
        fi
    done

    # Check essential files
    local essential_files=(
        "bsext_init"
        "uninstall.sh"
        "sh/init-extension"
        "sh/cleanup-extension"
        "sh/setup_python_env"
        "usr/bin/python3"
        "usr/lib/librknnrt.so"
    )

    for file in "${essential_files[@]}"; do
        if [[ ! -f "install/$file" ]]; then
            warn "Missing file: install/$file"
            ((errors++))
        fi
    done

    if [[ $errors -eq 0 ]]; then
        success "Structure verification passed"
    else
        warn "Structure verification found $errors issues"
    fi

    # Show package size
    local install_size=$(du -sh install/ 2>/dev/null | cut -f1)
    log "Extension package size: $install_size"
}

# Create development package
create_development_package() {
    log "Creating development package..."

    cd install
    zip -r "../$DEVELOPMENT_PACKAGE" ./ >/dev/null
    cd ..

    local package_size=$(du -sh "$DEVELOPMENT_PACKAGE" 2>/dev/null | cut -f1)
    success "Development package created: $DEVELOPMENT_PACKAGE ($package_size)"

    echo ""
    echo "Development Package Usage:"
    echo "1. Transfer $DEVELOPMENT_PACKAGE to player via DWS"
    echo "2. On player: mkdir -p /usr/local && cd /usr/local"
    echo "3. On player: unzip /storage/sd/$DEVELOPMENT_PACKAGE"
    echo "4. On player: source sh/pydev-env"
    echo "   Alternative: source sh/setup_python_env (auto-detects location)"
    echo "Note: Development installation is volatile (lost on reboot)"
}

# Create extension package
create_extension_package() {
    log "Creating production extension package..."

    cd install

    # Run make-extension script
    ../sh/make-extension-lvm || error "Extension creation failed"

    # Package the extension
    zip "../$EXTENSION_PACKAGE" ext_pydev* >/dev/null

    # Clean up temporary files
    rm -rf ext_pydev*

    cd ..

    local package_size=$(du -sh "$EXTENSION_PACKAGE" 2>/dev/null | cut -f1)
    success "Extension package created: $EXTENSION_PACKAGE ($package_size)"

    echo ""
    echo "Extension Package Usage:"
    echo "1. Transfer $EXTENSION_PACKAGE to player via DWS"
    echo "2. On player: mkdir -p /usr/local && cd /usr/local"
    echo "3. On player: unzip /storage/sd/$EXTENSION_PACKAGE"
    echo "4. On player: bash ./ext_pydev_install-lvm.sh"
    echo "5. On player: reboot"
    echo "Note: Extension installation is permanent (persists across reboots)"
}

# Run validation if requested
run_validation() {
    if [[ "$VERIFY" == "true" ]]; then
        echo ""
        log "Running validation..."
        if [[ -f "validate-build.sh" ]]; then
            ./validate-build.sh
        else
            warn "validate-build.sh not found, skipping validation"
        fi
    fi
}

# Main packaging function
main() {
    echo "BrightSign Python CV Extension - Package Creation"
    echo "================================================"

    local start_time=$(date +%s)

    check_prerequisites
    echo ""

    create_install_structure
    copy_sdk_components
    add_extension_scripts
    copy_rknn_wheel
    copy_user_init_examples
    copy_yolox_example
    echo ""

    verify_structure
    echo ""

    # Create packages based on options
    if [[ "$EXT_ONLY" == "true" ]]; then
        create_extension_package
    elif [[ "$DEV_ONLY" == "true" ]]; then
        create_development_package
    else
        create_development_package
        echo ""
        create_extension_package
    fi

    run_validation

    local end_time=$(date +%s)
    local duration=$((end_time - start_time))

    echo ""
    success "Packaging completed in $(($duration / 60))m $(($duration % 60))s"

    # Show created packages
    echo ""
    echo "Created packages:"
    if [[ "$EXT_ONLY" != "true" ]]; then
        echo "  📦 $DEVELOPMENT_PACKAGE (development/testing)"
    fi
    if [[ "$DEV_ONLY" != "true" ]]; then
        echo "  📦 $EXTENSION_PACKAGE (production)"
    fi

    echo ""
    echo "Next steps:"
    echo "1. Transfer package(s) to BrightSign player"
    echo "2. Install using instructions shown above"
    echo "3. Test with user init scripts from user-init/"
}

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