#!/bin/bash
#
# A2Z Radiology AI - Unified Installer
#
# Usage: curl -fsSL https://install.a2zradiology.ai | sh
#
# This script sets up A2Z Radiology AI in either:
#   - Cloud API mode: Access via HTTPS, no local installation
#   - On-Premise Container mode: Run inference locally with Docker
#

set -e

# Version
VERSION="1.2.0"

# =============================================================================
# CONFIGURATION
# =============================================================================
# URLs are injected at build time via CI/CD. Environment variables override
# for local development/testing.
#
# For non-interactive mode, also set:
#   A2Z_USERNAME      - Customer username
#   A2Z_PASSWORD      - Customer password
#   A2Z_NONINTERACTIVE=true

# Build-time placeholders (replaced by CI/CD)
# Format: __PLACEHOLDER__ gets replaced with actual value
BUILD_API_SERVER="https://staging.api.a2zradiology.ai"
BUILD_AUTH_SERVER="https://staging.auth.a2zradiology.ai"

# Resolve server URLs: env var > build-time value > error
resolve_server_url() {
    local env_value="$1"
    local build_value="$2"
    local var_name="$3"

    # Environment variable takes priority
    if [ -n "$env_value" ]; then
        echo "$env_value"
        return 0
    fi

    # Build-time value (if placeholder was replaced)
    if [ -n "$build_value" ] && [[ "$build_value" != __* ]]; then
        echo "$build_value"
        return 0
    fi

    # No value available - error
    echo "ERROR: $var_name not configured" >&2
    echo "" >&2
    echo "For local development, set environment variables:" >&2
    echo "  export A2Z_API_SERVER=https://staging.api.a2zradiology.ai" >&2
    echo "  export A2Z_AUTH_SERVER=https://staging.auth.a2zradiology.ai" >&2
    echo "" >&2
    return 1
}

# URL resolution is deferred until after argument parsing (so --help works without config)
SERVER_URL=""
AUTH_SERVER=""

# Release URL for container downloads (optional for cloud-only mode)
RELEASE_URL="${A2Z_RELEASE_URL:-}"

# Direct URLs (for presigned S3 URLs or custom hosting)
CONTAINER_URL="${A2Z_CONTAINER_URL:-}"
SIGNATURE_URL="${A2Z_SIGNATURE_URL:-}"
GPG_KEY_URL="${A2Z_GPG_KEY_URL:-}"
SAMPLES_URL="${A2Z_SAMPLES_URL:-}"

# Non-interactive mode (for CI/CD)
# Auto-detect if stdin is not a terminal (piped from curl)
if [ ! -t 0 ]; then
    PIPED_INPUT=true
else
    PIPED_INPUT=false
fi
NONINTERACTIVE="${A2Z_NONINTERACTIVE:-false}"

# Insecure mode (for testing with self-signed certs)
INSECURE="${A2Z_INSECURE:-false}"
CURL_OPTS=""
if [ "$INSECURE" = "true" ]; then
    CURL_OPTS="-k"
fi

# Installation directories
A2Z_HOME="${A2Z_HOME:-$HOME/a2z}"
A2Z_SAMPLES_DIR="$A2Z_HOME/samples"
A2Z_DOCS_DIR="$A2Z_HOME/docs"
A2Z_BIN_DIR="/usr/local/bin"

# Container image details
IMAGE_NAME="a2z-inference-a2z-unified-triage"
IMAGE_TAG="v${VERSION}"
IMAGE_FULL="${IMAGE_NAME}:${IMAGE_TAG}"

# Deployment mode (set via menu or --mode flag)
DEPLOYMENT_MODE=""

# Colors (disable in non-interactive mode for cleaner logs)
if [ "$NONINTERACTIVE" = "true" ]; then
    RED=''
    GREEN=''
    YELLOW=''
    BLUE=''
    CYAN=''
    BOLD=''
    NC=''
else
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    CYAN='\033[0;36m'
    BOLD='\033[1m'
    NC='\033[0m'
fi

# State
SKIP_SAMPLE_RUN=false

# Authentication state
AUTH_TOKEN=""
AUTH_USERNAME=""

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --mode)
            DEPLOYMENT_MODE="$2"
            if [ "$DEPLOYMENT_MODE" != "cloud" ] && [ "$DEPLOYMENT_MODE" != "container" ]; then
                echo "Invalid mode: $DEPLOYMENT_MODE (must be 'cloud' or 'container')"
                exit 1
            fi
            shift 2
            ;;
        --skip-sample-run)
            SKIP_SAMPLE_RUN=true
            shift
            ;;
        --test)
            # Force test run (for CI/CD)
            SKIP_SAMPLE_RUN=false
            shift
            ;;
        --help|-h)
            echo "A2Z Radiology AI Installer v$VERSION"
            echo ""
            echo "Usage: curl -fsSL https://install.a2zradiology.ai | sh"
            echo "       curl -fsSL https://install.a2zradiology.ai | sh -s -- [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --mode MODE          Deployment mode: 'cloud' or 'container' (skips menu)"
            echo "  --skip-sample-run    Skip running inference on sample data"
            echo "  --test               Force test run (for CI/CD)"
            echo "  --help, -h           Show this help message"
            echo ""
            echo "Environment variables (required):"
            echo "  A2Z_API_SERVER       API server URL"
            echo "  A2Z_AUTH_SERVER      Auth server URL"
            echo ""
            echo "Environment variables (for non-interactive mode):"
            echo "  A2Z_NONINTERACTIVE   Set to 'true' for CI/CD"
            echo "  A2Z_USERNAME         Username"
            echo "  A2Z_PASSWORD         Password"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# =============================================================================
# RESOLVE SERVER URLS (after argument parsing so --help works without config)
# =============================================================================
SERVER_URL=$(resolve_server_url "$A2Z_API_SERVER" "$BUILD_API_SERVER" "A2Z_API_SERVER") || exit 1
AUTH_SERVER=$(resolve_server_url "$A2Z_AUTH_SERVER" "$BUILD_AUTH_SERVER" "A2Z_AUTH_SERVER") || exit 1

# Print functions
print_banner() {
    echo ""
    echo -e "${CYAN}${BOLD}"
    echo "  ╔═══════════════════════════════════════════════════════════╗"
    echo "  ║                                                           ║"
    echo "  ║         A2Z Radiology AI Installer v${VERSION}              ║"
    echo "  ║                                                           ║"
    echo "  ╚═══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

print_step() {
    local step=$1
    local total=$2
    local message=$3
    echo ""
    echo -e "${BLUE}${BOLD}[$step/$total]${NC} ${BOLD}$message${NC}"
}

print_success() {
    echo -e "       ${GREEN}✓${NC} $1"
}

print_error() {
    echo -e "       ${RED}✗${NC} $1"
}

print_warning() {
    echo -e "       ${YELLOW}!${NC} $1"
}

print_info() {
    echo -e "       ${CYAN}→${NC} $1"
}

# Check if command exists
command_exists() {
    command -v "$1" &> /dev/null
}

# Progress bar for downloads
download_with_progress() {
    local url=$1
    local output=$2
    local description=$3

    if command_exists curl; then
        curl -fSL $CURL_OPTS --progress-bar -o "$output" "$url"
    elif command_exists wget; then
        wget --progress=bar:force -O "$output" "$url" 2>&1
    else
        print_error "Neither curl nor wget found"
        return 1
    fi
}

# State for features (set during authentication)
FEATURES=""
LICENSE_TYPE="standard"

# =============================================================================
# AUTHENTICATION (runs FIRST - fail fast on bad credentials)
# =============================================================================

check_basic_prerequisites() {
    # Minimal checks needed before authentication
    local all_ok=true

    # Check curl (required for API calls)
    if ! command_exists curl; then
        print_error "curl not found - required for API calls"
        all_ok=false
    fi

    # Check jq (required for JSON parsing)
    if ! command_exists jq; then
        print_error "jq not found - required for JSON parsing"
        print_info "Install with: sudo apt-get install jq (Linux) or brew install jq (macOS)"
        all_ok=false
    fi

    if [ "$all_ok" = false ]; then
        exit 1
    fi
}

authenticate() {
    print_step 1 3 "Authenticating..."

    # Check API server connectivity first
    print_info "Connecting to $SERVER_URL..."
    if ! curl -s $CURL_OPTS --connect-timeout 10 "$SERVER_URL/api/v1/health" > /dev/null 2>&1; then
        print_error "Cannot reach API server at $SERVER_URL"
        print_info "Check your network connection or A2Z_API_SERVER environment variable"
        exit 1
    fi

    # Try saved tokens first (skip in CI/non-interactive with explicit credentials)
    if [ "$NONINTERACTIVE" != "true" ] || [ -z "$A2Z_PASSWORD" ]; then
        if try_saved_tokens; then
            # Verify the token still works
            if curl -s $CURL_OPTS "$SERVER_URL/api/v1/health" \
                -H "Authorization: Bearer $AUTH_TOKEN" > /dev/null 2>&1; then
                print_success "Authenticated as $AUTH_USERNAME (saved session)"
                return 0
            fi
        fi
    fi

    # Need fresh login - get credentials
    local username=""
    local password=""

    if [ "$NONINTERACTIVE" = "true" ]; then
        username="${A2Z_USERNAME:-}"
        password="${A2Z_PASSWORD:-}"

        if [ -z "$username" ] || [ -z "$password" ]; then
            print_error "Non-interactive mode requires A2Z_USERNAME and A2Z_PASSWORD"
            exit 1
        fi
        print_info "Using credentials from environment"
    elif [ "$PIPED_INPUT" = "true" ]; then
        # Piped input without saved session - need env vars
        username="${A2Z_USERNAME:-}"
        password="${A2Z_PASSWORD:-}"

        if [ -z "$username" ] || [ -z "$password" ]; then
            print_error "No saved session and credentials not available"
            echo ""
            echo "  When running via pipe, provide credentials via environment:"
            echo ""
            echo "    A2Z_USERNAME=you@example.com A2Z_PASSWORD=xxx \\"
            echo "      curl -fsSL https://install.a2zradiology.ai | bash -s -- --mode cloud"
            echo ""
            exit 1
        fi
        print_info "Using credentials from environment"
    else
        echo ""
        echo -e "       ${BOLD}Enter your A2Z credentials:${NC}"
        echo ""
        read -p "       Username: " username
        read -sp "       Password: " password
        echo ""
    fi

    # Authenticate via API
    print_info "Validating credentials..."

    local auth_response
    auth_response=$(curl -s $CURL_OPTS -X POST "$SERVER_URL/api/v1/auth/login" \
        -H "Content-Type: application/json" \
        -d "{\"username\": \"$username\", \"password\": \"$password\"}" 2>/dev/null)

    # Check for access_token in response
    if ! echo "$auth_response" | jq -e '.access_token' > /dev/null 2>&1; then
        local error_msg
        error_msg=$(echo "$auth_response" | jq -r '.detail // .message // "Authentication failed"' 2>/dev/null)
        print_error "Authentication failed: $error_msg"
        exit 1
    fi

    # Extract tokens from response
    AUTH_TOKEN=$(echo "$auth_response" | jq -r '.access_token')
    AUTH_USERNAME="$username"

    local refresh_token
    local expires_in
    refresh_token=$(echo "$auth_response" | jq -r '.refresh_token')
    expires_in=$(echo "$auth_response" | jq -r '.expires_in // 3600')

    # Extract features/license type from response
    FEATURES=$(echo "$auth_response" | jq -r '.features // [] | join(",")' 2>/dev/null)
    if [[ "$FEATURES" == *"include_raw_predictions"* ]]; then
        LICENSE_TYPE="full"
        print_success "Authenticated as $username (license: regulatory/research)"
    else
        LICENSE_TYPE="standard"
        print_success "Authenticated as $username (license: standard)"
    fi

    # Store tokens (NOT password) for future sessions
    store_tokens "$username" "$AUTH_TOKEN" "$refresh_token" "$expires_in"
    print_success "Session saved (tokens stored, password not saved)"
}

store_tokens() {
    local username=$1
    local access_token=$2
    local refresh_token=$3
    local expires_in=$4

    mkdir -p "$A2Z_HOME/.config"
    chmod 700 "$A2Z_HOME/.config"

    # Calculate expiry timestamp
    local expires_at=$(($(date +%s) + expires_in))

    cat > "$A2Z_HOME/.config/credentials" << EOF
# A2Z Radiology AI - Token Storage
# Generated: $(date -Iseconds)
#
# These are OAuth tokens, NOT your password.
# Tokens are revocable from your A2Z dashboard.
# Delete this file to force re-authentication.

A2Z_USERNAME=$username
A2Z_ACCESS_TOKEN=$access_token
A2Z_REFRESH_TOKEN=$refresh_token
A2Z_TOKEN_EXPIRES_AT=$expires_at
A2Z_API_SERVER=$SERVER_URL
A2Z_AUTH_SERVER=$AUTH_SERVER
EOF
    chmod 600 "$A2Z_HOME/.config/credentials"
}

# Try to authenticate using saved tokens
try_saved_tokens() {
    local creds_file="$A2Z_HOME/.config/credentials"

    if [ ! -f "$creds_file" ]; then
        return 1
    fi

    # Load saved credentials
    source "$creds_file"

    # Check if we have required fields
    if [ -z "$A2Z_ACCESS_TOKEN" ] || [ -z "$A2Z_REFRESH_TOKEN" ]; then
        return 1
    fi

    # Check if access token is still valid (with 60s buffer)
    local now=$(date +%s)
    local expires_at="${A2Z_TOKEN_EXPIRES_AT:-0}"

    if [ "$now" -lt "$((expires_at - 60))" ]; then
        # Token still valid
        AUTH_TOKEN="$A2Z_ACCESS_TOKEN"
        AUTH_USERNAME="$A2Z_USERNAME"
        print_info "Using saved session for $A2Z_USERNAME"
        return 0
    fi

    # Token expired, try to refresh
    print_info "Session expired, refreshing..."
    if refresh_auth_token "$A2Z_REFRESH_TOKEN" "$A2Z_USERNAME"; then
        return 0
    fi

    # Refresh failed, need fresh login
    print_warning "Session expired, please log in again"
    rm -f "$creds_file"
    return 1
}

# Refresh access token using refresh token
refresh_auth_token() {
    local refresh_token=$1
    local username=$2

    local refresh_response
    refresh_response=$(curl -s $CURL_OPTS -X POST "$SERVER_URL/api/v1/auth/refresh" \
        -H "Authorization: Bearer $refresh_token" \
        -H "Content-Type: application/json" 2>/dev/null)

    # Check for new access_token
    if ! echo "$refresh_response" | jq -e '.access_token' > /dev/null 2>&1; then
        return 1
    fi

    # Extract new tokens
    local new_access_token
    local new_refresh_token
    local new_expires_in

    new_access_token=$(echo "$refresh_response" | jq -r '.access_token')
    new_refresh_token=$(echo "$refresh_response" | jq -r '.refresh_token')
    new_expires_in=$(echo "$refresh_response" | jq -r '.expires_in')

    # Update global state
    AUTH_TOKEN="$new_access_token"
    AUTH_USERNAME="$username"

    # Save new tokens
    store_tokens "$username" "$new_access_token" "$new_refresh_token" "$new_expires_in"

    print_success "Session refreshed"
    return 0
}

# =============================================================================
# MODE SELECTION (runs AFTER authentication)
# =============================================================================

ask_deployment_mode() {
    # If mode already set via --mode flag, use it
    if [ -n "$DEPLOYMENT_MODE" ]; then
        return
    fi

    # Non-interactive mode requires explicit --mode
    if [ "$NONINTERACTIVE" = "true" ]; then
        print_error "Non-interactive mode requires --mode flag (cloud or container)"
        exit 1
    fi

    # Piped input (curl | bash) requires --mode flag since we can't prompt
    if [ "$PIPED_INPUT" = "true" ]; then
        echo ""
        print_error "Interactive mode selection not available when piped"
        echo ""
        echo "  Please specify the deployment mode explicitly:"
        echo ""
        echo "    curl -fsSL https://install.a2zradiology.ai | bash -s -- --mode cloud"
        echo "    curl -fsSL https://install.a2zradiology.ai | bash -s -- --mode container"
        echo ""
        exit 1
    fi

    print_step 2 3 "Select deployment mode..."
    echo ""
    echo -e "    ${CYAN}[1]${NC} Cloud API"
    echo "        Access via HTTPS, no local installation"
    echo "        Data processed in A2Z secure cloud"
    echo ""
    echo -e "    ${CYAN}[2]${NC} On-Premise Container"
    echo "        Run inference locally with Docker"
    echo "        Data stays on your infrastructure"
    echo ""

    while true; do
        read -p "  Select mode [1/2]: " mode_choice
        case $mode_choice in
            1) DEPLOYMENT_MODE="cloud"; break ;;
            2) DEPLOYMENT_MODE="container"; break ;;
            *) echo "  Please enter 1 or 2" ;;
        esac
    done
}

# =============================================================================
# CLOUD API MODE FUNCTIONS
# =============================================================================

check_cloud_prerequisites() {
    # Cloud mode just needs curl and jq (already checked in basic prereqs)
    print_info "Cloud API mode - no additional prerequisites"
}

ask_test_setup() {
    if [ "$SKIP_SAMPLE_RUN" = true ]; then
        return 1
    fi

    if [ "$NONINTERACTIVE" = "true" ]; then
        # In non-interactive mode, run test by default unless skipped
        return 0
    fi

    # Piped input - skip test since we can't prompt
    if [ "$PIPED_INPUT" = "true" ]; then
        print_info "Skipping test prompt (piped input) - use --skip-sample-run to suppress"
        return 1
    fi

    echo ""
    echo -e "       ${BOLD}Would you like to verify your setup with a test run?${NC}"
    echo ""
    echo "         [1] Yes - run inference on a sample study (recommended)"
    echo "         [2] No - I'll test later"
    echo ""

    while true; do
        read -p "       Choice [1/2]: " test_choice
        case $test_choice in
            1) return 0 ;;
            2) return 1 ;;
            *) echo "       Please enter 1 or 2" ;;
        esac
    done
}

download_sample_for_cloud() {
    print_info "Downloading sample study..."

    mkdir -p "$A2Z_HOME/samples"
    cd "$A2Z_HOME"

    # Get sample download URL from API
    print_info "Fetching sample data URL..."

    local download_response
    download_response=$(curl -s $CURL_OPTS -X GET "$SERVER_URL/api/v1/releases/download-urls" \
        -H "Authorization: Bearer $AUTH_TOKEN" \
        -H "Content-Type: application/json" 2>/dev/null)

    local samples_url
    samples_url=$(echo "$download_response" | jq -r '.samples_url // empty' 2>/dev/null)

    if [ -z "$samples_url" ]; then
        print_warning "Could not get sample URL from server"
        print_info "Skipping sample test - you can test with your own DICOM data"
        return 1
    fi

    # Download samples
    print_info "Downloading sample DICOM study..."
    if download_with_progress "$samples_url" "samples.tar.gz" "Samples"; then
        tar -xzf samples.tar.gz -C "$A2Z_HOME/samples" 2>/dev/null || {
            print_warning "Failed to extract samples"
            return 1
        }
        rm -f samples.tar.gz
        print_success "Sample data downloaded to $A2Z_HOME/samples"
        return 0
    else
        print_warning "Failed to download samples"
        return 1
    fi
}

run_cloud_inference() {
    print_info "Running test inference..."

    # Find sample DICOM directory
    local sample_dir
    sample_dir=$(find "$A2Z_HOME/samples" -type d -name "positive_case" 2>/dev/null | head -1)
    if [ -z "$sample_dir" ]; then
        sample_dir=$(find "$A2Z_HOME/samples" -type d -name "dicom" 2>/dev/null | head -1)
    fi
    if [ -z "$sample_dir" ]; then
        sample_dir=$(find "$A2Z_HOME/samples" -type d -mindepth 1 -maxdepth 2 2>/dev/null | head -1)
    fi

    if [ -z "$sample_dir" ] || [ ! -d "$sample_dir" ]; then
        print_warning "Sample DICOM directory not found"
        return 1
    fi

    print_info "Uploading study from: $sample_dir"

    # Find DICOM files
    local dicom_files
    dicom_files=$(find "$sample_dir" -type f \( -name "*.dcm" -o ! -name "*.*" \) 2>/dev/null | head -50)

    if [ -z "$dicom_files" ]; then
        print_warning "No DICOM files found in sample directory"
        return 1
    fi

    # Build multipart form data for upload
    local upload_args=()
    while IFS= read -r file; do
        upload_args+=(-F "files=@$file")
    done <<< "$dicom_files"

    # Upload to API
    print_info "Uploading DICOM files..."
    local upload_response
    upload_response=$(curl -s $CURL_OPTS -X POST "$SERVER_URL/api/v1/upload" \
        -H "Authorization: Bearer $AUTH_TOKEN" \
        "${upload_args[@]}" 2>/dev/null)

    local job_id
    job_id=$(echo "$upload_response" | jq -r '.job_id // .id // empty' 2>/dev/null)

    if [ -z "$job_id" ]; then
        local error_msg
        error_msg=$(echo "$upload_response" | jq -r '.detail // "Upload failed"' 2>/dev/null)
        print_error "Upload failed: $error_msg"
        return 1
    fi

    print_success "Upload complete - Job ID: $job_id"

    # Poll for status
    print_info "Processing study..."
    local max_attempts=60
    local attempt=0
    local status=""

    while [ $attempt -lt $max_attempts ]; do
        attempt=$((attempt + 1))

        local status_response
        status_response=$(curl -s $CURL_OPTS "$SERVER_URL/api/v1/inference/$job_id/status" \
            -H "Authorization: Bearer $AUTH_TOKEN" 2>/dev/null)

        status=$(echo "$status_response" | jq -r '.status // "unknown"' 2>/dev/null)

        # Show progress
        local progress=$((attempt * 100 / max_attempts))
        printf "\r       Processing... [%-20s] %d%%" "$(printf '#%.0s' $(seq 1 $((progress / 5))))" "$progress"

        if [ "$status" = "completed" ] || [ "$status" = "COMPLETED" ]; then
            echo ""
            print_success "Processing complete!"
            break
        elif [ "$status" = "failed" ] || [ "$status" = "FAILED" ]; then
            echo ""
            local error_msg
            error_msg=$(echo "$status_response" | jq -r '.error // "Unknown error"' 2>/dev/null)
            print_error "Processing failed: $error_msg"
            return 1
        fi

        sleep 5
    done

    if [ "$status" != "completed" ] && [ "$status" != "COMPLETED" ]; then
        echo ""
        print_error "Timeout waiting for results"
        return 1
    fi

    # Get results
    local results_response
    results_response=$(curl -s $CURL_OPTS "$SERVER_URL/api/v1/inference/$job_id/results" \
        -H "Authorization: Bearer $AUTH_TOKEN" 2>/dev/null)

    # Save results
    mkdir -p "$A2Z_HOME/test-results"
    echo "$results_response" | jq '.' > "$A2Z_HOME/test-results/sample-result.json" 2>/dev/null

    # Display results
    display_cloud_results "$results_response"

    return 0
}

display_cloud_results() {
    local results=$1

    echo ""
    echo -e "  ${BOLD}═══════════════════════ RESULTS ═══════════════════════${NC}"
    echo ""

    # Check if predictions exist
    if ! echo "$results" | jq -e '.predictions' > /dev/null 2>&1; then
        print_warning "No predictions in response"
        echo "$results" | jq '.' 2>/dev/null
        return
    fi

    # Header
    printf "  ${BOLD}%-30s │ %-8s │ %-10s${NC}\n" "Condition" "Status" "Confidence"
    echo "  ──────────────────────────────┼──────────┼────────────"

    # Parse and display each prediction
    echo "$results" | jq -r '.predictions | to_entries[] | "\(.key)|\(.value.predicted_class // "unknown")|\(.value.raw_prediction // 0)"' 2>/dev/null | while IFS='|' read -r condition status confidence; do
        # Format condition name
        local display_condition
        display_condition=$(echo "$condition" | sed 's/_/ /g')

        # Format status with color
        local status_display
        if [ "$status" = "present" ] || [ "$status" = "PRESENT" ]; then
            status_display="${YELLOW}PRESENT${NC}"
        else
            status_display="${GREEN}ABSENT${NC}"
        fi

        # Format confidence as percentage
        local conf_pct
        conf_pct=$(echo "$confidence" | awk '{printf "%.1f%%", $1 * 100}')

        printf "  %-30s │ %-17s │ %s\n" "$display_condition" "$status_display" "$conf_pct"
    done

    echo ""
    echo -e "  ${GREEN}Results saved to:${NC} $A2Z_HOME/test-results/sample-result.json"
    echo ""
}

print_cloud_complete() {
    echo ""
    echo -e "${GREEN}${BOLD}"
    echo "  ╔═══════════════════════════════════════════════════════════╗"
    echo "  ║                                                           ║"
    echo "  ║              Cloud API Setup Complete!                    ║"
    echo "  ║                                                           ║"
    echo "  ╚═══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo ""
    echo -e "  ${BOLD}API Endpoint:${NC}"
    echo "    $SERVER_URL"
    echo ""
    echo -e "  ${BOLD}Quick Reference - Upload a study:${NC}"
    echo ""
    echo "    # 1. Login and get token"
    echo "    TOKEN=\$(curl -s -X POST $SERVER_URL/api/v1/auth/login \\"
    echo "      -H 'Content-Type: application/json' \\"
    echo "      -d '{\"username\":\"YOUR_USERNAME\",\"password\":\"YOUR_PASSWORD\"}' | jq -r '.access_token')"
    echo ""
    echo "    # 2. Upload DICOM files"
    echo "    curl -X POST $SERVER_URL/api/v1/upload \\"
    echo "      -H \"Authorization: Bearer \$TOKEN\" \\"
    echo "      -F 'files=@/path/to/study/*.dcm'"
    echo ""
    echo "    # 3. Check status and get results"
    echo "    curl $SERVER_URL/api/v1/inference/{job_id}/results \\"
    echo "      -H \"Authorization: Bearer \$TOKEN\""
    echo ""
    echo -e "  ${BOLD}Documentation:${NC}"
    echo "    https://docs.a2zradiology.ai"
    echo ""
    echo -e "  ${BOLD}Support:${NC}"
    echo "    support@a2zradiology.ai"
    echo ""
}

setup_cloud_api() {
    check_cloud_prerequisites

    # Authentication already done in main()
    # Now offer to test the setup
    if ask_test_setup; then
        if download_sample_for_cloud; then
            run_cloud_inference || true
        fi
    fi

    print_cloud_complete
}

# =============================================================================
# CONTAINER MODE FUNCTIONS
# =============================================================================

# Step 1: Check prerequisites (for container mode)
check_prerequisites() {
    print_step 1 7 "Checking prerequisites..."

    local all_ok=true

    # Check Docker
    if command_exists docker; then
        local docker_version=$(docker --version 2>/dev/null | cut -d' ' -f3 | tr -d ',')
        print_success "Docker $docker_version installed"

        # Check if Docker daemon is running
        if docker info > /dev/null 2>&1; then
            print_success "Docker daemon running"
        else
            print_error "Docker daemon not running"
            print_info "Start Docker with: sudo systemctl start docker"
            all_ok=false
        fi
    else
        print_error "Docker not found"
        print_info "Install Docker: https://docs.docker.com/get-docker/"
        all_ok=false
    fi

    # Check GPG
    if command_exists gpg; then
        print_success "GPG installed"
    else
        print_error "GPG not found"
        print_info "Install with: sudo apt-get install gnupg"
        all_ok=false
    fi

    # Check Python 3
    if command_exists python3; then
        local python_version=$(python3 --version 2>&1 | cut -d' ' -f2)
        print_success "Python $python_version installed"
    else
        print_error "Python 3 not found"
        all_ok=false
    fi

    # Check disk space (need at least 10GB free)
    if command_exists df; then
        local available_gb=$(df -BG "$HOME" 2>/dev/null | tail -1 | awk '{print $4}' | tr -d 'G')
        if [ -n "$available_gb" ] && [ "$available_gb" -ge 10 ]; then
            print_success "${available_gb}GB disk space available"
        else
            print_warning "Low disk space: ${available_gb}GB (recommend 10GB+)"
        fi
    fi

    # Check network connectivity
    if curl -s $CURL_OPTS --connect-timeout 5 "$AUTH_SERVER_DEFAULT" > /dev/null 2>&1; then
        print_success "Network connectivity OK"
    else
        print_warning "Cannot reach auth server - check network"
    fi

    if [ "$all_ok" = false ]; then
        echo ""
        print_error "Prerequisites not met. Please install missing dependencies."
        exit 1
    fi
}

# Step 2: Authenticate (BEFORE downloading anything)
authenticate_user() {
    print_step 2 7 "Authenticating..."

    local username=""
    local password=""
    local auth_server=""

    # Non-interactive mode: use environment variables
    if [ "$NONINTERACTIVE" = "true" ]; then
        username="${A2Z_USERNAME:-}"
        password="${A2Z_PASSWORD:-}"
        auth_server="${A2Z_AUTH_SERVER:-$AUTH_SERVER_DEFAULT}"

        if [ -z "$username" ] || [ -z "$password" ]; then
            print_error "Non-interactive mode requires A2Z_USERNAME and A2Z_PASSWORD environment variables"
            exit 1
        fi
        print_info "Using credentials from environment variables"
    else
        # Interactive mode: prompt for credentials
        echo ""
        echo -e "       ${BOLD}Enter your A2Z credentials:${NC}"
        echo ""

        # Get username
        read -p "       Username: " username

        # Get password (hidden)
        read -sp "       Password: " password
        echo ""

        # Auth server (use default)
        auth_server="$AUTH_SERVER_DEFAULT"
    fi

    # Authenticate with auth server
    print_info "Validating credentials..."

    local auth_response
    local http_code

    # Make auth request and capture both response body and HTTP status code
    auth_response=$(curl -s $CURL_OPTS -w "\n%{http_code}" -X POST "$auth_server/login" \
        -H "Content-Type: application/json" \
        -d "{\"username\": \"$username\", \"password\": \"$password\"}" 2>/dev/null)

    # Extract HTTP status code (last line) and response body
    http_code=$(echo "$auth_response" | tail -n1)
    auth_response=$(echo "$auth_response" | sed '$d')

    # Debug output in non-interactive mode
    if [ "$NONINTERACTIVE" = "true" ]; then
        print_info "Auth server: $auth_server"
        print_info "HTTP status: $http_code"
        print_info "Response: $auth_response"
    fi

    # Check if authentication was successful
    local success=$(echo "$auth_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('success', False))" 2>/dev/null)

    if [ "$success" != "True" ]; then
        local message=$(echo "$auth_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('message', 'Authentication failed'))" 2>/dev/null)
        print_error "Authentication failed: $message"
        exit 1
    fi

    print_success "Authentication successful"

    # Extract access token for subsequent API calls
    AUTH_TOKEN=$(echo "$auth_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('access_token', ''))" 2>/dev/null)

    if [ -z "$AUTH_TOKEN" ]; then
        print_error "Failed to extract access token from authentication response"
        exit 1
    fi

    # Extract features from response to determine license type
    FEATURES=$(echo "$auth_response" | python3 -c "import sys,json; d=json.load(sys.stdin); f=d.get('features',[]); print(','.join(f) if f else '')" 2>/dev/null)

    if [[ "$FEATURES" == *"include_raw_predictions"* ]]; then
        LICENSE_TYPE="full"
        print_success "License: regulatory/research (includes probability scores)"
    else
        LICENSE_TYPE="standard"
        print_success "License: standard"
    fi

    # Store credentials in keyring or fallback
    store_credentials "$username" "$password" "$auth_server"
}

# Store credentials securely
store_credentials() {
    local username=$1
    local password=$2
    local auth_server=$3

    # Check if keyring is available
    if ! python3 -c "import keyring" 2>/dev/null; then
        print_info "Installing keyring package..."
        pip3 install --user keyring secretstorage 2>/dev/null || {
            sudo pip3 install keyring secretstorage 2>/dev/null || true
        }
    fi

    # Try to store in keyring
    if python3 -c "import keyring" 2>/dev/null; then
        python3 -c "
import keyring
keyring.set_password('a2z-radiology', 'username', '''$username''')
keyring.set_password('a2z-radiology', 'password', '''$password''')
keyring.set_password('a2z-radiology', 'auth_server', '''$auth_server''')
" 2>/dev/null && {
            print_success "Credentials stored in OS keyring"
            return 0
        }
    fi

    # Fallback: store in file
    print_warning "Keyring unavailable - using fallback storage"
    mkdir -p "$A2Z_HOME/.config"
    cat > "$A2Z_HOME/.config/credentials" << EOF
A2Z_USERNAME=$username
A2Z_PASSWORD=$password
A2Z_AUTH_SERVER=$auth_server
EOF
    chmod 600 "$A2Z_HOME/.config/credentials"
    print_success "Credentials saved to $A2Z_HOME/.config/credentials"
}

# Step 2.5: Get download URLs from device-server
# This replaces hardcoded RELEASE_URL with dynamic presigned URLs
get_download_urls() {
    # Skip if URLs already provided via environment variables
    if [ -n "$CONTAINER_URL" ] && [ -n "$SIGNATURE_URL" ] && [ -n "$GPG_KEY_URL" ]; then
        print_info "Using provided download URLs (skipping server lookup)"
        return 0
    fi

    print_step "2.5" 7 "Fetching download URLs..."

    if [ -z "$AUTH_TOKEN" ]; then
        print_error "No authentication token - authenticate first"
        exit 1
    fi

    # Call device-server to get presigned download URLs
    local download_response
    download_response=$(curl -s $CURL_OPTS -X GET "$SERVER_URL/releases/download-urls" \
        -H "Authorization: Bearer $AUTH_TOKEN" \
        -H "Content-Type: application/json" 2>/dev/null)

    # Check if request was successful
    local success=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('success', False))" 2>/dev/null)

    if [ "$success" != "True" ]; then
        local detail=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('detail', 'Failed to get download URLs'))" 2>/dev/null)
        print_error "Failed to get download URLs: $detail"
        exit 1
    fi

    # Extract URLs from response
    CONTAINER_URL=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('container_url', ''))" 2>/dev/null)
    SIGNATURE_URL=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('signature_url', ''))" 2>/dev/null)
    GPG_KEY_URL=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('gpg_key_url', ''))" 2>/dev/null)
    SAMPLES_URL=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('samples_url', ''))" 2>/dev/null)

    local release_version=$(echo "$download_response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('version', 'unknown'))" 2>/dev/null)

    if [ -z "$CONTAINER_URL" ] || [ -z "$SIGNATURE_URL" ] || [ -z "$GPG_KEY_URL" ]; then
        print_error "Incomplete download URLs received from server"
        exit 1
    fi

    print_success "Download URLs retrieved (release $release_version)"
}

# Step 3: Download container image
download_container() {
    print_step 3 7 "Downloading container image (~2GB)..."

    mkdir -p "$A2Z_HOME/downloads"
    cd "$A2Z_HOME/downloads"

    # Use direct URLs if provided, otherwise use RELEASE_URL
    local container_src="${CONTAINER_URL:-${RELEASE_URL}/container.tar.gz}"
    local signature_src="${SIGNATURE_URL:-${RELEASE_URL}/container.tar.gz.asc}"
    local gpg_key_src="${GPG_KEY_URL:-${RELEASE_URL}/gpg-public-key.asc}"

    # Detect file extension from URL (handle .tar vs .tar.gz)
    local container_ext="tar.gz"
    if [[ "$container_src" == *.tar ]] && [[ "$container_src" != *.tar.gz ]]; then
        container_ext="tar"
    fi

    # Download container
    print_info "Downloading container image..."
    download_with_progress "$container_src" "container.${container_ext}" "Container"

    # Download signature (matches container extension)
    print_info "Downloading signature..."
    download_with_progress "$signature_src" "container.${container_ext}.asc" "Signature"

    # Download GPG key
    print_info "Downloading GPG public key..."
    download_with_progress "$gpg_key_src" "gpg-public-key.asc" "GPG Key"

    print_success "Downloads complete"
}

# Step 4: Verify GPG signature
verify_signature() {
    print_step 4 7 "Verifying GPG signature..."

    cd "$A2Z_HOME/downloads"

    # Import GPG key
    print_info "Importing A2Z public key..."
    gpg --import gpg-public-key.asc 2>/dev/null

    # Find the container file and its signature
    local container_file=""
    local signature_file=""
    if [ -f "container.tar.gz" ]; then
        container_file="container.tar.gz"
        signature_file="container.tar.gz.asc"
    elif [ -f "container.tar" ]; then
        container_file="container.tar"
        signature_file="container.tar.asc"
    else
        print_error "Container file not found for verification"
        exit 1
    fi

    # Verify signature
    print_info "Verifying container signature..."
    if gpg --verify "$signature_file" "$container_file" 2>/dev/null; then
        print_success "Signature valid - container is authentic"
    else
        print_error "SIGNATURE VERIFICATION FAILED"
        print_error "Container may have been tampered with!"
        print_error "Contact support@a2zradiology.ai immediately"
        exit 1
    fi
}

# Step 5: Load Docker image
load_image() {
    print_step 5 7 "Loading Docker image..."

    cd "$A2Z_HOME/downloads"

    # Find the container file (could be .tar or .tar.gz)
    local container_file=""
    if [ -f "container.tar.gz" ]; then
        print_info "Decompressing image..."
        gunzip -f container.tar.gz 2>/dev/null || true
        container_file="container.tar"
    elif [ -f "container.tar" ]; then
        container_file="container.tar"
    else
        print_error "Container file not found"
        exit 1
    fi

    # Load into Docker
    print_info "Loading into Docker (this may take a minute)..."
    if docker load -i "$container_file" 2>/dev/null; then
        print_success "Image loaded: $IMAGE_FULL"
    else
        print_error "Failed to load Docker image"
        exit 1
    fi

    # Cleanup download
    rm -f container.tar container.tar.gz container.tar.gz.asc gpg-public-key.asc
    print_success "Cleaned up temporary files"
}

# Step 6: Download samples (license-specific)
download_samples() {
    print_step 6 7 "Downloading samples (for $LICENSE_TYPE license)..."

    mkdir -p "$A2Z_SAMPLES_DIR"
    mkdir -p "$A2Z_DOCS_DIR"

    cd "$A2Z_HOME/downloads"

    # Select sample package based on license type
    local samples_src
    if [ -n "$SAMPLES_URL" ]; then
        # Direct URL provided
        samples_src="$SAMPLES_URL"
        print_info "Using provided samples URL"
    elif [ "$LICENSE_TYPE" = "full" ]; then
        samples_src="${RELEASE_URL}/samples-full.tar.gz"
        print_info "Downloading full samples (with probability scores)"
    else
        samples_src="${RELEASE_URL}/samples-standard.tar.gz"
        print_info "Downloading standard samples"
    fi

    # Download samples
    print_info "Downloading sample DICOM studies..."
    download_with_progress "$samples_src" "samples.tar.gz" "Samples"

    # Extract samples
    print_info "Extracting samples..."
    tar -xzf samples.tar.gz -C "$A2Z_SAMPLES_DIR"
    rm -f samples.tar.gz
    print_success "Samples extracted to $A2Z_SAMPLES_DIR"

    # Download docs
    print_info "Downloading documentation..."
    download_with_progress "${RELEASE_URL}/docs.tar.gz" "docs.tar.gz" "Docs"
    tar -xzf docs.tar.gz -C "$A2Z_DOCS_DIR"
    rm -f docs.tar.gz
    print_success "Documentation installed"

    # Run sample inference if not skipped
    if [ "$SKIP_SAMPLE_RUN" = false ]; then
        print_info "Running inference on sample study..."

        # Get credentials
        local creds_file=""
        if [ -f "$A2Z_HOME/.config/credentials" ]; then
            creds_file="$A2Z_HOME/.config/credentials"
        fi

        # Find a sample study
        local sample_dir=$(find "$A2Z_SAMPLES_DIR" -type d -name "positive_case" 2>/dev/null | head -1)
        if [ -z "$sample_dir" ]; then
            sample_dir=$(find "$A2Z_SAMPLES_DIR" -type d -name "dicom" 2>/dev/null | head -1)
        fi

        if [ -n "$sample_dir" ] && [ -d "$sample_dir" ]; then
            local output_file="$A2Z_SAMPLES_DIR/sample_result.json"

            # Run inference
            if [ -n "$creds_file" ]; then
                docker run --rm \
                    --cpus="4.0" --memory="8g" \
                    --env-file "$creds_file" \
                    -v "$sample_dir":/input:ro \
                    -v "$A2Z_SAMPLES_DIR":/outputs \
                    "$IMAGE_FULL" \
                    python src/lean_wrapper.py \
                        --directory_path /input \
                        --output_path /outputs/sample_result.json 2>/dev/null && {
                    print_success "Sample inference complete!"
                    print_success "Output: $output_file"
                } || {
                    print_warning "Sample inference failed - check credentials"
                }
            else
                print_warning "Skipping sample run - credentials not in fallback file"
            fi
        else
            print_warning "Sample DICOM not found - skipping test run"
        fi
    else
        print_info "Skipping sample run (--skip-sample-run)"
    fi
}

# Install CLI wrapper
install_cli_wrapper() {
    print_info "Installing a2z-inference command..."

    # Create wrapper script (version is injected from installer)
    local wrapper_script
    wrapper_script=$(cat << 'WRAPPER_EOF'
#!/bin/bash
#
# A2Z Inference CLI
# Usage: a2z-inference /path/to/study -o output.json
#

A2Z_HOME="${A2Z_HOME:-$HOME/a2z}"
VERSION="__VERSION__"
IMAGE="${A2Z_IMAGE:-a2z-inference-a2z-unified-triage:v${VERSION}}"
INSTALLER_URL="${A2Z_INSTALLER_URL:-https://install.a2zradiology.ai}"

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

# Exit codes
EXIT_OK=0
EXIT_MISSING_CREDS=1
EXIT_NETWORK_ERROR=2
EXIT_AUTH_FAILED=3
EXIT_INVALID_INPUT=4
EXIT_INFERENCE_ERROR=5

# Print functions
print_error() {
    echo -e "${RED}${BOLD}ERROR:${NC} $1" >&2
}

print_success() {
    echo -e "${GREEN}${BOLD}SUCCESS:${NC} $1"
}

print_info() {
    echo -e "${BLUE}[*]${NC} $1"
}

print_step() {
    echo -e "${BOLD}==> $1${NC}"
}

# Helper: Get credentials from keyring or fallback file
get_credential() {
    local key=$1
    # Try keyring first
    local val=$(python3 -c "import keyring; v=keyring.get_password('a2z-radiology', '$key'); print(v or '')" 2>/dev/null)
    if [ -n "$val" ]; then
        echo "$val"
        return
    fi
    # Fallback to file
    if [ -f "$A2Z_HOME/.config/credentials" ]; then
        grep "^A2Z_$(echo $key | tr 'a-z' 'A-Z')=" "$A2Z_HOME/.config/credentials" | cut -d= -f2-
    fi
}

# Pre-flight auth check
check_authentication() {
    local username="$1"
    local password="$2"
    local auth_server="$3"

    print_step "Authenticating..."

    local response
    local http_code

    response=$(curl -s -w "\n%{http_code}" -X POST "$auth_server/login" \
        -H "Content-Type: application/json" \
        -d "{\"username\": \"$username\", \"password\": \"$password\"}" 2>&1)

    http_code=$(echo "$response" | tail -n1)
    response=$(echo "$response" | sed '$d')

    # Network error
    if [ -z "$http_code" ] || [ "$http_code" = "000" ]; then
        echo ""
        print_error "Cannot connect to authentication server"
        echo ""
        echo -e "  ${BOLD}Server:${NC} $auth_server"
        echo -e "  ${BOLD}Action:${NC} Check your internet connection"
        echo ""
        exit $EXIT_NETWORK_ERROR
    fi

    # Check for success
    local success=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('success', False))" 2>/dev/null)

    if [ "$success" != "True" ]; then
        local message=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('message', 'Unknown error'))" 2>/dev/null)
        echo ""
        print_error "Authentication failed"
        echo ""
        echo -e "  ${BOLD}Reason:${NC} $message"
        echo -e "  ${BOLD}User:${NC}   $username"
        echo ""
        echo -e "  ${BOLD}Possible fixes:${NC}"
        echo "    - Check your username and password"
        echo "    - Run the installer again to update credentials"
        echo "    - Contact support@a2zradiology.ai"
        echo ""
        exit $EXIT_AUTH_FAILED
    fi

    echo -e "  ${GREEN}Authenticated as${NC} $username"
}

# Run inference
run_inference() {
    local input_dir="$1"
    local output_file="$2"
    local username="$3"
    local password="$4"
    local auth_server="$5"

    print_step "Running inference..."
    echo -e "  ${BOLD}Input:${NC}  $input_dir"
    echo -e "  ${BOLD}Output:${NC} $output_file"
    echo ""

    # Auto-detect host UID for permission compatibility
    local host_uid=$(id -u)
    local host_gid=$(id -g)

    # Run container
    docker run --rm \
        --user "$host_uid:$host_gid" \
        --cpus="4.0" --memory="16g" \
        -e A2Z_USERNAME="$username" \
        -e A2Z_PASSWORD="$password" \
        -e A2Z_AUTH_SERVER="$auth_server" \
        -v "$input_dir":/input:ro \
        -v "$output_dir":/outputs \
        "$IMAGE" \
        python src/lean_wrapper.py \
            --directory_path /input \
            --output_path "/outputs/$output_name" 2>&1

    local exit_code=$?
    return $exit_code
}

# Main
main() {
    local INPUT_DIR=""
    local OUTPUT_FILE=""

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -o|--output)
                OUTPUT_FILE="$2"
                shift 2
                ;;
            --help|-h)
                echo "A2Z Inference CLI v$VERSION"
                echo ""
                echo "Usage: a2z-inference /path/to/study -o output.json"
                echo ""
                echo "Options:"
                echo "  -o, --output FILE  Output JSON file path (default: ./result.json)"
                echo "  --version          Show version information"
                echo "  --update           Update to latest version"
                echo "  -h, --help         Show this help"
                echo ""
                echo "Exit codes:"
                echo "  0 - Success"
                echo "  1 - Missing credentials"
                echo "  2 - Network error"
                echo "  3 - Authentication failed"
                echo "  4 - Invalid input"
                echo "  5 - Inference error"
                echo ""
                echo "Examples:"
                echo "  a2z-inference /path/to/dicom/study"
                echo "  a2z-inference /path/to/dicom/study -o results.json"
                echo ""
                exit 0
                ;;
            --version)
                echo "a2z-inference v$VERSION"
                echo "Container: $IMAGE"
                exit 0
                ;;
            --update)
                echo "Updating A2Z Inference..."
                if command -v curl &> /dev/null; then
                    curl -fsSL "$INSTALLER_URL" | sh
                elif command -v wget &> /dev/null; then
                    wget -qO- "$INSTALLER_URL" | sh
                else
                    print_error "curl or wget required for update"
                    exit 1
                fi
                exit $?
                ;;
            -*)
                print_error "Unknown option: $1"
                echo "Run 'a2z-inference --help' for usage"
                exit $EXIT_INVALID_INPUT
                ;;
            *)
                INPUT_DIR="$1"
                shift
                ;;
        esac
    done

    # Validate input
    if [ -z "$INPUT_DIR" ]; then
        print_error "No input directory specified"
        echo ""
        echo "Usage: a2z-inference /path/to/dicom/study -o output.json"
        exit $EXIT_INVALID_INPUT
    fi

    if [ ! -d "$INPUT_DIR" ]; then
        print_error "Input directory does not exist: $INPUT_DIR"
        exit $EXIT_INVALID_INPUT
    fi

    # Set defaults
    OUTPUT_FILE="${OUTPUT_FILE:-./result.json}"

    # Resolve paths
    INPUT_DIR=$(realpath "$INPUT_DIR")
    local output_dir=$(dirname "$(realpath -m "$OUTPUT_FILE")")
    local output_name=$(basename "$OUTPUT_FILE")

    mkdir -p "$output_dir"

    # Get credentials
    local username=$(get_credential username)
    local password=$(get_credential password)
    local auth_server=$(get_credential auth_server)
    auth_server="${auth_server:-https://device-auth-staging-831731287.us-east-1.elb.amazonaws.com}"

    if [ -z "$username" ] || [ -z "$password" ]; then
        echo ""
        print_error "Credentials not found"
        echo ""
        echo -e "  ${BOLD}Action:${NC} Run the installer again to set up credentials"
        echo ""
        echo "    curl -fsSL https://install.a2zradiology.ai | sh"
        echo ""
        exit $EXIT_MISSING_CREDS
    fi

    echo ""
    echo -e "${BOLD}A2Z Inference v$VERSION${NC}"
    echo ""

    # Step 1: Pre-flight auth check
    check_authentication "$username" "$password" "$auth_server"

    # Step 2: Run inference
    echo ""
    run_inference "$INPUT_DIR" "$OUTPUT_FILE" "$username" "$password" "$auth_server"

    local exit_code=$?
    if [ $exit_code -ne 0 ]; then
        echo ""
        print_error "Inference failed (exit code $exit_code)"
        echo ""
        echo -e "  ${BOLD}Possible causes:${NC}"
        echo "    - Invalid DICOM data"
        echo "    - Unsupported image orientation (must be axial)"
        echo "    - Insufficient system resources"
        echo ""
        echo -e "  ${BOLD}Need help?${NC} support@a2zradiology.ai"
        echo ""
        exit $EXIT_INFERENCE_ERROR
    fi

    echo ""
    print_success "Inference complete"
    echo ""
    echo -e "  ${BOLD}Output:${NC} $OUTPUT_FILE"
    echo ""
}

main "$@"
WRAPPER_EOF
)

    # Replace version placeholder
    wrapper_script="${wrapper_script/__VERSION__/$VERSION}"

    # Try to install to /usr/local/bin
    if [ -w "$A2Z_BIN_DIR" ] || sudo -n true 2>/dev/null; then
        echo "$wrapper_script" | sudo tee "$A2Z_BIN_DIR/a2z-inference" > /dev/null
        sudo chmod +x "$A2Z_BIN_DIR/a2z-inference"
        print_success "Installed: a2z-inference"
    else
        # Install to user's local bin
        local user_bin="$HOME/.local/bin"
        mkdir -p "$user_bin"
        echo "$wrapper_script" > "$user_bin/a2z-inference"
        chmod +x "$user_bin/a2z-inference"
        print_success "Installed: $user_bin/a2z-inference"
        print_warning "Add to PATH: export PATH=\"\$HOME/.local/bin:\$PATH\""
    fi
}

# Print completion message
print_complete() {
    echo ""
    echo -e "${GREEN}${BOLD}"
    echo "  ╔═══════════════════════════════════════════════════════════╗"
    echo "  ║                                                           ║"
    echo "  ║              Installation Complete!                       ║"
    echo "  ║                                                           ║"
    echo "  ╚═══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo ""
    echo -e "  ${BOLD}Run inference:${NC}"
    echo "    a2z-inference /path/to/dicom/study -o result.json"
    echo ""
    echo -e "  ${BOLD}Documentation:${NC}"
    echo "    $A2Z_DOCS_DIR/QUICK_REFERENCE.md"
    echo ""
    echo -e "  ${BOLD}Sample output:${NC}"
    echo "    $A2Z_SAMPLES_DIR/sample_result.json"
    echo ""
    echo -e "  ${BOLD}Support:${NC}"
    echo "    support@a2zradiology.ai"
    echo ""
}

# Container mode setup (wraps existing flow)
setup_container() {
    # Check container-specific prerequisites (Docker, GPG, etc.)
    check_prerequisites

    # Authentication already done in main(), just get download URLs
    get_download_urls        # Get presigned URLs from device-server
    download_container       # Download container
    verify_signature         # Verify GPG signature
    load_image               # Load Docker image
    download_samples         # Download license-specific samples
    install_cli_wrapper
    print_complete
}

# =============================================================================
# MAIN ENTRY POINT
# =============================================================================
# Flow:
#   1. Banner + basic prereqs check
#   2. Authenticate (FIRST - fail fast)
#   3. Mode selection (now that we know user is valid)
#   4. Mode-specific setup

main() {
    print_banner

    # Check basic prerequisites (curl, jq)
    check_basic_prerequisites

    # Authenticate FIRST - fail fast on bad credentials
    authenticate

    # Now select deployment mode
    ask_deployment_mode

    # Run mode-specific setup
    print_step 3 3 "Setting up $DEPLOYMENT_MODE mode..."

    case "$DEPLOYMENT_MODE" in
        cloud)
            setup_cloud_api
            ;;
        container)
            setup_container
            ;;
        *)
            print_error "Unknown deployment mode: $DEPLOYMENT_MODE"
            exit 1
            ;;
    esac
}

main "$@"
