#!/bin/bash
# Script untuk menjalankan Locust Load Test dengan virtual environment

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="$SCRIPT_DIR/venv"
PYTHON="$VENV_DIR/bin/python"
PIP="$VENV_DIR/bin/pip"
LOCUST="$VENV_DIR/bin/locust"

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

echo "🚀 Elemes LMS Load Test Setup"
echo "=============================="

# Check if virtual environment exists
if [ ! -d "$VENV_DIR" ]; then
    echo -e "${YELLOW}📦 Creating virtual environment...${NC}"
    python3 -m venv "$VENV_DIR"
    echo -e "${GREEN}✅ Virtual environment created${NC}"
else
    echo -e "${GREEN}✅ Virtual environment already exists${NC}"
fi

# Activate virtual environment and install dependencies
echo -e "${YELLOW}📦 Installing dependencies...${NC}"
"$PIP" install --upgrade pip -q
"$PIP" install -r "$SCRIPT_DIR/requirements.txt" -q
echo -e "${GREEN}✅ Dependencies installed${NC}"

echo ""
echo "=============================="
echo "Usage:"
echo "  ./run.sh              - Start Web UI (default: API on :5000)"
echo "  ./run.sh frontend     - Start Web UI for Frontend testing (:3000)"
echo "  ./run.sh fullstack    - Start Web UI testing full stack (:3000)"
echo "  ./run.sh headless     - Run in headless mode (default: 50 users, 5 spawn rate, 60s)"
echo "  ./run.sh clean        - Remove virtual environment"
echo ""
echo "Headless mode options:"
echo "  ./run.sh headless -u 100 -r 10 -t 300s"
echo "  ./run.sh headless frontend -u 100 -r 10 -t 5m"
echo ""
echo "Examples:"
echo "  # Test API only (Flask backend)"
echo "  ./run.sh --host=http://localhost:5000"
echo ""
echo "  # Test Frontend only (SvelteKit)"
echo "  ./run.sh frontend --host=http://localhost:3000"
echo ""
echo "  # Test full stack (Frontend + API via proxy)"
echo "  ./run.sh fullstack -u 100 -r 10 -t 5m"
echo "=============================="
echo ""

# Parse arguments
MODE="api"
HEADLESS=false

if [ "$1" == "clean" ]; then
    echo -e "${YELLOW}🗑️  Removing virtual environment...${NC}"
    rm -rf "$VENV_DIR"
    echo -e "${GREEN}✅ Virtual environment removed${NC}"
    exit 0
elif [ "$1" == "frontend" ]; then
    MODE="frontend"
    shift
elif [ "$1" == "fullstack" ]; then
    MODE="fullstack"
    shift
elif [ "$1" == "headless" ]; then
    HEADLESS=true
    shift
    # Check if next arg is frontend/fullstack
    if [ "$1" == "frontend" ]; then
        MODE="frontend"
        shift
    elif [ "$1" == "fullstack" ]; then
        MODE="fullstack"
        shift
    fi
fi

# Set default host based on mode
if [ "$MODE" == "frontend" ]; then
    HOST="${HOST:-http://localhost:3000}"
    echo -e "${YELLOW}🎨 Frontend testing mode${NC}"
elif [ "$MODE" == "fullstack" ]; then
    HOST="${HOST:-http://localhost:3000}"
    echo -e "${YELLOW}🔄 Fullstack testing mode${NC}"
else
    HOST="${HOST:-http://localhost:5000}"
    echo -e "${YELLOW}🔧 API testing mode${NC}"
fi

if [ "$HEADLESS" == "true" ]; then
    # Default values (HOST already set above based on mode)
    USERS="${USERS:-50}"
    SPAWN="${SPAWN:-5}"
    TIME="${TIME:-60s}"
    
    # Parse additional arguments
    shift
    while [[ $# -gt 0 ]]; do
        case $1 in
            -u|--users)
                USERS="$2"
                shift 2
                ;;
            -r|--spawn)
                SPAWN="$2"
                shift 2
                ;;
            -t|--time)
                TIME="$2"
                shift 2
                ;;
            --host)
                HOST="$2"
                shift 2
                ;;
            *)
                echo -e "${RED}Unknown option: $1${NC}"
                exit 1
                ;;
        esac
    done
    
    echo -e "${YELLOW}🏃 Running Locust in headless mode...${NC}"
    echo "   Host: $HOST"
    echo "   Users: $USERS"
    echo "   Spawn rate: $SPAWN/s"
    echo "   Duration: $TIME"
    echo ""
    
    "$LOCUST" -f "$SCRIPT_DIR/locustfile.py" \
        --host="$HOST" \
        --headless \
        -u "$USERS" \
        -r "$SPAWN" \
        -t "$TIME"
else
    # Default: Web UI mode (HOST already set above based on mode)
    
    echo -e "${YELLOW}🌐 Starting Locust Web UI...${NC}"
    echo "   Target: $HOST"
    echo "   Mode: $MODE"
    echo "   Open http://localhost:8089 in your browser"
    echo ""
    echo -e "${GREEN}Press Ctrl+C to stop${NC}"
    echo ""
    
    "$LOCUST" -f "$SCRIPT_DIR/locustfile.py" --host="$HOST"
fi