#!/bin/bash

# Optimization script for multi-core, low-latency processing
# Sets environment variables and starts servers with optimal configurations

echo "=========================================="
echo "UI Analysis System - Optimized Setup"
echo "=========================================="

# Get number of CPU cores
NUM_CORES=$(nproc)
echo "Detected CPU cores: $NUM_CORES"

# Export optimization flags
export OMP_NUM_THREADS=$NUM_CORES
export MKL_NUM_THREADS=$NUM_CORES
export NUMEXPR_NUM_THREADS=$NUM_CORES
export OPENBLAS_NUM_THREADS=$NUM_CORES
export VECLIB_MAXIMUM_THREADS=$NUM_CORES
export TORCH_NUM_THREADS=$NUM_CORES
export PYTORCH_ENABLE_MPS_FALLBACK=1

# Enable optimal PyTorch settings
export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:512"

# Disable CPU throttling for latency
export PYTHONUNBUFFERED=1

echo "Environment variables set for $NUM_CORES cores"

# Calculate optimal worker count - use all cores for async workers
WORKERS=$NUM_CORES
echo "Starting servers with $WORKERS workers..."

echo ""
echo "[1] Starting OmniParser Server on port 8000..."
cd /workspaces/omoi-v2/OmniParser
python -m omnitool.omniparserserver.omniparserserver \
    --som_model_path weights/icon_detect/model.pt \
    --device cpu \
    --BOX_TRESHOLD 0.05 \
    --save_cropped_images \
    --cropped_images_dir /tmp/omoi_cropped \
    --host 0.0.0.0 \
    --port 8000 &

OMONIPARSER_PID=$!
echo "OmniParser Server PID: $OMNIPARSER_PID"

# Wait for OmniParser to start
sleep 5

echo ""
echo "[2] Starting UI Element API Server on port 8001..."
cd /workspaces/omoi-v2
python ui_element_api_server.py \
    --port 8001 \
    --workers $WORKERS \
    --host 0.0.0.0 &

UI_API_PID=$!
echo "UI Element API Server PID: $UI_API_PID"

# Wait for startup
sleep 3

echo ""
echo "=========================================="
echo "✓ All servers started!"
echo "=========================================="
echo ""
echo "Services:"
echo "  OmniParser: http://127.0.0.1:8000/probe/"
echo "  UI Element API: http://127.0.0.1:8001/docs"
echo ""
echo "Environment optimizations:"
echo "  - Multi-threading enabled ($NUM_CORES threads)"
echo "  - PyTorch optimizations active"
echo "  - $WORKERS async workers for API"
echo "  - Temp cropped images cache: /tmp/omoi_cropped"
echo ""
echo "Press Ctrl+C to stop all servers"
echo "=========================================="

# Wait for both processes
wait $OMNIPARSER_PID $UI_API_PID