#!/bin/bash
#
# Test script for scx_mitosis CPU cell isolation
#
# This script tests that processes in different managed child cgroups are
# scheduled on distinct CPU sets.
#
# Usage:
#   ./test_cell_isolation.sh [--num-cells N] [--workers N] [--test-dynamic] [--test-borrowing]
#                            [--test-enqueue-borrowing]
#                            [--test-rebalancing] [--test-cpuset-change] [--test-all]
#
# Options:
#   --num-cells N   Number of test cells to create (default: 2)
#   --workers N     Number of stress-ng workers per cell (default: 5)
#   --test-dynamic  Run dynamic cell creation/destruction tests
#   --test-rebalancing Run demand-based CPU rebalancing test
#   --test-cpuset-change Run cpuset change detection test
#   --test-new-cell-demand Run new cell demand seeding test
#   --test-all      Run all tests (basic + dynamic + borrowing + rebalancing + cpuset-change + new-cell-demand)
#   --help          Show this help message
#

set -e

# Configuration
SCHEDULER_BIN="${SCHEDULER_BIN:-./target/release/scx_mitosis}"
CGROUP_BASE="/sys/fs/cgroup/test.slice"
LOG_FILE="/tmp/scx_mitosis_test.log"

# Defaults
NUM_CELLS=2
WORKERS_PER_CELL=5
SAMPLE_COUNT=5
SAMPLE_INTERVAL=1
TEST_DYNAMIC=0
TEST_BORROWING=0
TEST_ENQUEUE_BORROWING=0
TEST_REBALANCING=0
TEST_CPUSET_CHANGE=0
TEST_NEW_CELL_DEMAND=0
TEST_ALL=0

# Test result tracking
declare -A TEST_RESULTS

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

usage() {
    head -21 "$0" | tail -18
    exit 0
}

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --num-cells)
            NUM_CELLS="$2"
            shift 2
            ;;
        --workers)
            WORKERS_PER_CELL="$2"
            shift 2
            ;;
        --test-dynamic)
            TEST_DYNAMIC=1
            shift
            ;;
        --test-borrowing)
            TEST_BORROWING=1
            shift
            ;;
        --test-enqueue-borrowing)
            TEST_ENQUEUE_BORROWING=1
            shift
            ;;
        --test-rebalancing)
            TEST_REBALANCING=1
            shift
            ;;
        --test-cpuset-change)
            TEST_CPUSET_CHANGE=1
            shift
            ;;
        --test-new-cell-demand)
            TEST_NEW_CELL_DEMAND=1
            shift
            ;;
        --test-all)
            TEST_ALL=1
            shift
            ;;
        --help)
            usage
            ;;
        *)
            log_error "Unknown option: $1"
            usage
            ;;
    esac
done

# Cleanup function
cleanup() {
    log_info "Cleaning up..."

    # Kill stress-ng processes
    sudo pkill -9 stress-ng 2>/dev/null || true

    # Kill scheduler (use pkill to ensure child processes are killed too)
    sudo pkill -9 scx_mitosis 2>/dev/null || true

    # Also try killing by PID if we have it
    if [[ -n "$SCHED_PID" ]] && ps -p "$SCHED_PID" > /dev/null 2>&1; then
        sudo kill -9 "$SCHED_PID" 2>/dev/null || true
    fi

    # Remove test cgroups (numbered cells)
    for i in $(seq 1 $NUM_CELLS); do
        cell_name="test_cell_${i}"
        cell_path="$CGROUP_BASE/$cell_name"
        if [[ -d "$cell_path" ]]; then
            # Kill any remaining processes
            for pid in $(cat "$cell_path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$cell_path" 2>/dev/null || true
        fi
    done

    # Remove dynamic test cgroups
    for cell_name in test_cell_dynamic test_cell_temp test_cell_3 test_cell_reuse_a test_cell_reuse_b test_cell_borrow_busy test_cell_borrow_idle test_cell_rebal_busy test_cell_rebal_idle test_cell_cpuset_a test_cell_cpuset_b test_cell_demand_a test_cell_demand_b test_cell_demand_c; do
        cell_path="$CGROUP_BASE/$cell_name"
        if [[ -d "$cell_path" ]]; then
            for pid in $(cat "$cell_path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$cell_path" 2>/dev/null || true
        fi
    done

    # Remove the parent test.slice cgroup
    if [[ -d "$CGROUP_BASE" ]]; then
        # Kill any remaining processes
        for pid in $(cat "$CGROUP_BASE/cgroup.procs" 2>/dev/null); do
            sudo kill -9 "$pid" 2>/dev/null || true
        done
        sudo rmdir "$CGROUP_BASE" 2>/dev/null || true
    fi

    wait 2>/dev/null || true
}

trap cleanup EXIT

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

    if [[ ! -x "$SCHEDULER_BIN" ]]; then
        log_error "Scheduler binary not found: $SCHEDULER_BIN"
        log_error "Please build with: cargo build --release -p scx_mitosis"
        exit 1
    fi

    if ! command -v stress-ng &> /dev/null; then
        log_error "stress-ng not found. Please install it."
        exit 1
    fi

    if [[ ! -d "/sys/fs/cgroup" ]]; then
        log_error "cgroup2 filesystem not found at /sys/fs/cgroup"
        exit 1
    fi

    # Check if we can read sched_ext state
    if [[ ! -f "/sys/kernel/sched_ext/state" ]]; then
        log_error "sched_ext not available (missing /sys/kernel/sched_ext/state)"
        exit 1
    fi
}

# Create test cgroups
create_cgroups() {
    log_info "Creating test cgroups under $CGROUP_BASE..."

    # Create parent if needed
    if [[ ! -d "$CGROUP_BASE" ]]; then
        sudo mkdir -p "$CGROUP_BASE"
    fi

    # Enable CPU controller for children (required for sched_ext cgroup callbacks)
    if ! grep -q "cpu" "$CGROUP_BASE/cgroup.subtree_control" 2>/dev/null; then
        log_info "Enabling CPU controller for $CGROUP_BASE children"
        echo "+cpu" | sudo tee "$CGROUP_BASE/cgroup.subtree_control" > /dev/null
    fi

    for i in $(seq 1 $NUM_CELLS); do
        cell_name="test_cell_${i}"
        cell_path="$CGROUP_BASE/$cell_name"

        if [[ ! -d "$cell_path" ]]; then
            sudo mkdir "$cell_path"
            log_info "Created $cell_path"
        fi
    done

    # List created cgroups
    ls -la "$CGROUP_BASE" | grep test_cell
}

# Start the scheduler
start_scheduler() {
    log_info "Starting scx_mitosis scheduler (extra flags: $@)..."

    sudo "$SCHEDULER_BIN" \
        --cell-parent-cgroup /test.slice \
        "$@" \
        > "$LOG_FILE" 2>&1 &
    SCHED_PID=$!

    log_info "Scheduler PID: $SCHED_PID"

    # Wait for scheduler to attach
    sleep 3

    # Verify scheduler is running
    if ! ps -p "$SCHED_PID" > /dev/null 2>&1; then
        log_error "Scheduler failed to start. Log output:"
        cat "$LOG_FILE"
        exit 1
    fi

    # Check sched_ext state
    local state=$(cat /sys/kernel/sched_ext/state 2>/dev/null)
    if [[ "$state" != "enabled" ]]; then
        log_error "sched_ext not enabled. State: $state"
        cat "$LOG_FILE"
        exit 1
    fi

    log_info "sched_ext state: $state"

    # Show which cells were created
    log_info "Scheduler log (cell creation):"
    grep -E "(Created cell|Assigning)" "$LOG_FILE" | head -10 || true
}

# Start workloads in each cell
start_workloads() {
    log_info "Starting $WORKERS_PER_CELL stress-ng workers in each cell..."

    for i in $(seq 1 $NUM_CELLS); do
        cell_name="test_cell_${i}"
        cell_path="$CGROUP_BASE/$cell_name"

        # Start stress-ng in a subshell that first moves to the cgroup
        sudo bash -c "echo \$\$ > $cell_path/cgroup.procs && \
            exec stress-ng --cpu $WORKERS_PER_CELL --timeout 120s --quiet" &

        log_info "Started workload in $cell_name"
    done

    # Wait for workers to start
    sleep 2
}

# Get CPU for a process
get_process_cpu() {
    local pid=$1
    if [[ -f "/proc/$pid/stat" ]]; then
        # Field 39 in /proc/pid/stat is the processor number
        awk '{print $39}' "/proc/$pid/stat" 2>/dev/null
    fi
}

# Sample CPU assignments
sample_cpu_assignments() {
    log_info "Sampling CPU assignments ($SAMPLE_COUNT samples)..."
    echo ""

    declare -A cell_cpus

    for sample in $(seq 1 $SAMPLE_COUNT); do
        echo "=== Sample $sample ==="

        for i in $(seq 1 $NUM_CELLS); do
            cell_name="test_cell_${i}"
            cell_path="$CGROUP_BASE/$cell_name"

            echo "  $cell_name:"

            for pid in $(cat "$cell_path/cgroup.procs" 2>/dev/null); do
                if [[ -d "/proc/$pid" ]]; then
                    cpu=$(get_process_cpu "$pid")
                    comm=$(cat "/proc/$pid/comm" 2>/dev/null)
                    if [[ -n "$cpu" ]]; then
                        echo "    PID $pid ($comm) -> CPU $cpu"
                        # Track unique CPUs per cell
                        cell_cpus["${cell_name}_${cpu}"]=1
                    fi
                fi
            done
        done

        if [[ $sample -lt $SAMPLE_COUNT ]]; then
            sleep $SAMPLE_INTERVAL
        fi
        echo ""
    done

    # Print summary of CPUs used by each cell
    echo "=== CPU Usage Summary ==="
    for i in $(seq 1 $NUM_CELLS); do
        cell_name="test_cell_${i}"
        cpus=""
        for key in "${!cell_cpus[@]}"; do
            if [[ "$key" == "${cell_name}_"* ]]; then
                cpu="${key#${cell_name}_}"
                if [[ -n "$cpus" ]]; then
                    cpus="$cpus, $cpu"
                else
                    cpus="$cpu"
                fi
            fi
        done
        echo "  $cell_name used CPUs: $cpus"
    done
}

# Check for CPU isolation (detect overlapping CPU usage)
check_isolation() {
    log_info "Checking CPU isolation..."

    declare -A cpu_to_cell
    local overlap_detected=0

    for i in $(seq 1 $NUM_CELLS); do
        cell_name="test_cell_${i}"
        cell_path="$CGROUP_BASE/$cell_name"

        for pid in $(cat "$cell_path/cgroup.procs" 2>/dev/null); do
            if [[ -d "/proc/$pid" ]]; then
                cpu=$(get_process_cpu "$pid")
                if [[ -n "$cpu" ]]; then
                    if [[ -n "${cpu_to_cell[$cpu]}" ]] && [[ "${cpu_to_cell[$cpu]}" != "$cell_name" ]]; then
                        log_warn "CPU $cpu used by both ${cpu_to_cell[$cpu]} and $cell_name"
                        overlap_detected=1
                    fi
                    cpu_to_cell[$cpu]="$cell_name"
                fi
            fi
        done
    done

    if [[ $overlap_detected -eq 0 ]]; then
        log_info "No CPU overlap detected between cells"
        return 0
    else
        log_warn "CPU overlap detected - cells are not fully isolated"
        return 1
    fi
}

# Record test result
record_result() {
    local test_name="$1"
    local result="$2"
    TEST_RESULTS["$test_name"]="$result"
    if [[ "$result" == "PASSED" ]]; then
        log_info "$test_name: ${GREEN}PASSED${NC}"
    elif [[ "$result" == "INCONCLUSIVE" ]]; then
        log_warn "$test_name: ${YELLOW}INCONCLUSIVE${NC}"
    else
        log_error "$test_name: ${RED}FAILED${NC}"
    fi
}

# Test: Basic CPU isolation
test_basic_isolation() {
    log_info "=== Running test: Basic CPU Isolation ==="

    start_workloads
    sample_cpu_assignments

    if check_isolation; then
        record_result "basic_isolation" "PASSED"
        return 0
    else
        record_result "basic_isolation" "INCONCLUSIVE"
        return 1
    fi
}

# Test: Dynamic cell lifecycle (creation, workload, destruction)
test_dynamic_cell_lifecycle() {
    log_info "=== Running test: Dynamic Cell Lifecycle ==="

    local cell_path="$CGROUP_BASE/test_cell_dynamic"

    # Cleanup any existing cgroup
    if [[ -d "$cell_path" ]]; then
        sudo rmdir "$cell_path" 2>/dev/null || true
        sleep 1
    fi

    # Create a new cgroup while scheduler is running
    sudo mkdir "$cell_path"
    sleep 3  # Wait for inotify to fire and scheduler to process

    # Verify scheduler log shows new cell
    if ! grep -q "Created cell.*test_cell_dynamic" "$LOG_FILE"; then
        log_error "Dynamic cell creation not detected in logs"
        record_result "dynamic_cell_lifecycle" "FAILED"
        return 1
    fi

    log_info "Cell creation detected in logs"

    # Start workload in new cell
    sudo bash -c "echo \$\$ > $cell_path/cgroup.procs && \
        exec stress-ng --cpu 2 --timeout 60s --quiet" &
    local workload_pid=$!
    sleep 2

    # Verify tasks are running
    if ! pgrep -f "stress-ng.*cpu" > /dev/null; then
        log_error "Workload not running in dynamic cell"
        record_result "dynamic_cell_lifecycle" "FAILED"
        return 1
    fi

    # Check that processes in the dynamic cell are running
    local procs=$(cat "$cell_path/cgroup.procs" 2>/dev/null | wc -l)
    if [[ "$procs" -lt 1 ]]; then
        log_error "No processes found in dynamic cell cgroup"
        record_result "dynamic_cell_lifecycle" "FAILED"
        return 1
    fi

    log_info "Found $procs processes in dynamic cell"

    # Stop the workload
    sudo pkill -f "stress-ng.*cpu" 2>/dev/null || true
    sleep 2

    # Remove the cell
    sudo rmdir "$cell_path"
    sleep 3

    # Verify destruction logged
    if ! grep -q "Destroyed cell.*test_cell_dynamic" "$LOG_FILE"; then
        log_error "Cell destruction not detected in logs"
        log_info "Log contents related to test_cell_dynamic:"
        grep "test_cell_dynamic" "$LOG_FILE" || echo "(no matches)"
        record_result "dynamic_cell_lifecycle" "FAILED"
        return 1
    fi

    log_info "Cell destruction detected in logs"
    record_result "dynamic_cell_lifecycle" "PASSED"
    return 0
}

# Test: CPU reallocation after adding a cell
test_cpu_reallocation() {
    log_info "=== Running test: CPU Reallocation ==="

    # Count lines with cell configuration messages before adding new cell
    local config_count_before=$(grep -c "Cell config updated\|Applied initial cell configuration" "$LOG_FILE" || echo "0")

    # Add a third cell
    local cell_path="$CGROUP_BASE/test_cell_3"
    if [[ -d "$cell_path" ]]; then
        sudo rmdir "$cell_path" 2>/dev/null || true
        sleep 1
    fi

    sudo mkdir "$cell_path"
    sleep 3

    # Count configuration applications after
    local config_count_after=$(grep -c "Cell config updated\|Applied initial cell configuration" "$LOG_FILE" || echo "0")

    if [[ "$config_count_after" -gt "$config_count_before" ]]; then
        log_info "Cell configuration was reapplied after adding new cell"
        log_info "Config applications: $config_count_before -> $config_count_after"
        record_result "cpu_reallocation" "PASSED"
        return 0
    else
        log_error "CPU reallocation not triggered after adding cell"
        record_result "cpu_reallocation" "FAILED"
        return 1
    fi
}

# Test: Cell ID reuse
test_cell_id_reuse() {
    log_info "=== Running test: Cell ID Reuse ==="

    local cell_path_a="$CGROUP_BASE/test_cell_reuse_a"
    local cell_path_b="$CGROUP_BASE/test_cell_reuse_b"

    # Clean up any existing cells
    for path in "$cell_path_a" "$cell_path_b"; do
        if [[ -d "$path" ]]; then
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done
    sleep 1

    # Create first cell
    sudo mkdir "$cell_path_a"
    sleep 2

    # Get the cell ID from log
    local cell_id=$(grep "Created cell.*test_cell_reuse_a" "$LOG_FILE" | tail -1 | grep -oP "cell \K\d+" || echo "")

    if [[ -z "$cell_id" ]]; then
        log_error "Could not find cell ID for test_cell_reuse_a"
        record_result "cell_id_reuse" "FAILED"
        return 1
    fi

    log_info "First cell created with ID: $cell_id"

    # Destroy it
    sudo rmdir "$cell_path_a"
    sleep 2

    # Create a new cell - should reuse the same ID
    sudo mkdir "$cell_path_b"
    sleep 2

    local new_cell_id=$(grep "Created cell.*test_cell_reuse_b" "$LOG_FILE" | tail -1 | grep -oP "cell \K\d+" || echo "")

    if [[ -z "$new_cell_id" ]]; then
        log_error "Could not find cell ID for test_cell_reuse_b"
        record_result "cell_id_reuse" "FAILED"
        return 1
    fi

    log_info "Second cell created with ID: $new_cell_id"

    if [[ "$cell_id" == "$new_cell_id" ]]; then
        log_info "Cell ID was reused (ID $cell_id)"
        record_result "cell_id_reuse" "PASSED"
        return 0
    else
        # ID reuse is not guaranteed if other cells exist - this is acceptable
        log_warn "Cell ID not reused (expected $cell_id, got $new_cell_id)"
        log_info "This may be acceptable depending on other active cells"
        record_result "cell_id_reuse" "INCONCLUSIVE"
        return 0
    fi
}

# Test: CPU borrowing - verify that a busy cell borrows CPUs from idle cells.
# Args:
#   $1 - stressor type: "pipe" (wakeup-heavy, select_cpu path) or
#                        "cpu" (CPU spinners, enqueue path)
#   $2 - test name for result recording
test_borrowing() {
    local stressor="$1"
    local test_name="$2"
    log_info "=== Running test: CPU Borrowing ($stressor stressor) ==="

    # Kill existing scheduler if running
    if [[ -n "$SCHED_PID" ]] && ps -p "$SCHED_PID" > /dev/null 2>&1; then
        sudo kill "$SCHED_PID" 2>/dev/null || true
        sleep 2
    fi
    sudo pkill -9 scx_mitosis 2>/dev/null || true
    sleep 2

    # Create two test cells
    local busy_path="$CGROUP_BASE/test_cell_borrow_busy"
    local idle_path="$CGROUP_BASE/test_cell_borrow_idle"

    for path in "$busy_path" "$idle_path"; do
        if [[ -d "$path" ]]; then
            for pid in $(cat "$path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done

    sudo mkdir -p "$busy_path"
    sudo mkdir -p "$idle_path"

    # Start scheduler with borrowing
    start_scheduler --enable-borrowing

    # Use more workers than the busy cell has CPUs to force borrowing
    local workers=$(nproc)

    log_info "Starting $workers stress-ng $stressor workers in busy cell..."
    sudo bash -c "echo \$\$ > $busy_path/cgroup.procs && \
        exec stress-ng --$stressor $workers --timeout 60s --quiet" &

    # Wait for stats to accumulate
    sleep 5

    # Capture monitor output
    local monitor_output="/tmp/scx_mitosis_monitor_borrow.json"
    log_info "Capturing monitor output..."
    timeout 4 "$SCHEDULER_BIN" --monitor 1 > "$monitor_output" 2>/dev/null || true

    # Kill workloads
    sudo pkill -9 stress-ng 2>/dev/null || true
    sleep 1

    # Validate: check borrowing and usage tracking metrics from monitor JSON
    local result
    result=$(python3 -c "
import json, sys

with open('$monitor_output', 'r') as f:
    text = f.read()

decoder = json.JSONDecoder()
last_obj = None
pos = 0
while pos < len(text):
    while pos < len(text) and text[pos] in ' \t\n\r':
        pos += 1
    if pos >= len(text):
        break
    try:
        obj, end = decoder.raw_decode(text, pos)
        last_obj = obj
        pos = end
    except json.JSONDecodeError:
        break

if last_obj is None:
    print('NO_DATA')
    sys.exit(0)

borrowed_pct = last_obj.get('borrowed_pct', 0)
cells = last_obj.get('cells', {})

if borrowed_pct <= 0:
    print('FAIL:borrowed_pct is 0 (CSTAT_BORROWED not incremented)')
    sys.exit(0)

# Find busy cell (highest util) and idle cell (lowest util)
busy_id = max(cells, key=lambda c: cells[c].get('util_pct', 0))
idle_id = min(cells, key=lambda c: cells[c].get('util_pct', 0))
busy = cells[busy_id]
idle = cells[idle_id]

if busy.get('util_pct', 0) <= 0:
    print('FAIL:no cell has util_pct > 0')
    sys.exit(0)

if busy.get('demand_borrow_pct', 0) <= 0:
    print('FAIL:busy cell %s has demand_borrow_pct=0 (should be borrowing)' % busy_id)
    sys.exit(0)

if idle.get('lent_pct', 0) <= 0:
    print('FAIL:idle cell %s has lent_pct=0 (should be lending)' % idle_id)
    sys.exit(0)

msg = 'PASS:borrowed=%.2f,busy=%s(util=%.1f,borrow=%.1f),idle=%s(lent=%.1f)' % (
    borrowed_pct, busy_id, busy['util_pct'], busy['demand_borrow_pct'],
    idle_id, idle['lent_pct'])
print(msg)
" 2>&1)

    log_info "Borrowing ($stressor) result: $result"

    case "$result" in
        PASS:*)
            record_result "$test_name" "PASSED"
            return 0
            ;;
        FAIL:*)
            local reason="${result#FAIL:}"
            log_error "Borrowing test ($stressor) failed: $reason"
            record_result "$test_name" "FAILED"
            return 1
            ;;
        NO_DATA)
            log_error "Borrowing test ($stressor): no monitor data captured"
            record_result "$test_name" "FAILED"
            return 1
            ;;
        *)
            log_error "Borrowing test ($stressor): unexpected result: $result"
            record_result "$test_name" "FAILED"
            return 1
            ;;
    esac
}

# Test: Demand-based CPU rebalancing - busy cell should get more CPUs
test_cpu_rebalancing_demand() {
    log_info "=== Running test: Demand-Based CPU Rebalancing ==="

    # Kill existing scheduler if running
    if [[ -n "$SCHED_PID" ]] && ps -p "$SCHED_PID" > /dev/null 2>&1; then
        sudo kill "$SCHED_PID" 2>/dev/null || true
        sleep 2
    fi
    sudo pkill -9 scx_mitosis 2>/dev/null || true
    sleep 2

    # Create two test cells
    local busy_path="$CGROUP_BASE/test_cell_rebal_busy"
    local idle_path="$CGROUP_BASE/test_cell_rebal_idle"

    for path in "$busy_path" "$idle_path"; do
        if [[ -d "$path" ]]; then
            for pid in $(cat "$path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done

    sudo mkdir -p "$busy_path"
    sudo mkdir -p "$idle_path"

    # Start scheduler with borrowing and rebalancing enabled
    log_info "Starting scx_mitosis with --enable-borrowing --enable-rebalancing..."
    sudo "$SCHEDULER_BIN" \
        --cell-parent-cgroup /test.slice \
        --enable-borrowing \
        --enable-rebalancing \
        --rebalance-cooldown-s 3 \
        --rebalance-threshold 10 \
        > "$LOG_FILE" 2>&1 &
    SCHED_PID=$!

    sleep 3

    # Verify scheduler is running
    if ! ps -p "$SCHED_PID" > /dev/null 2>&1; then
        log_error "Scheduler failed to start. Log output:"
        cat "$LOG_FILE"
        record_result "cpu_rebalancing_demand" "FAILED"
        return 1
    fi

    local state=$(cat /sys/kernel/sched_ext/state 2>/dev/null)
    if [[ "$state" != "enabled" ]]; then
        log_error "sched_ext not enabled. State: $state"
        cat "$LOG_FILE"
        record_result "cpu_rebalancing_demand" "FAILED"
        return 1
    fi

    # Capture baseline CPU counts before workload (equal distribution)
    local baseline_output="/tmp/scx_mitosis_monitor_baseline.json"
    log_info "Capturing baseline CPU counts (before workload)..."
    timeout 4 "$SCHEDULER_BIN" --monitor 1 > "$baseline_output" 2>/dev/null || true

    # Start heavy workload only in busy cell
    local workers=$(nproc)
    log_info "Starting $workers stress-ng pipe workers in busy cell..."
    sudo bash -c "echo \$\$ > $busy_path/cgroup.procs && \
        exec stress-ng --pipe $workers --timeout 60s --quiet" &

    # Wait for demand stats and rebalancing to trigger
    # Need ~5s warmup for EWMA + 3s cooldown
    log_info "Waiting for rebalancing to trigger..."
    sleep 10

    # Capture monitor output while workload is still running
    local monitor_output="/tmp/scx_mitosis_monitor_rebal.json"
    log_info "Capturing monitor output..."
    timeout 4 "$SCHEDULER_BIN" --monitor 1 > "$monitor_output" 2>/dev/null || true

    # Kill stress-ng workers
    sudo pkill -9 stress-ng 2>/dev/null || true
    sleep 1

    # Validate using the scheduler log (cell-to-name mapping) and monitor output (CPU counts)
    local result
    result=$(python3 -c "
import json, sys, re

def parse_last_json(path):
    try:
        with open(path, 'r') as f:
            text = f.read()
    except FileNotFoundError:
        return None
    decoder = json.JSONDecoder()
    last_obj = None
    pos = 0
    while pos < len(text):
        while pos < len(text) and text[pos] in ' \t\n\r':
            pos += 1
        if pos >= len(text):
            break
        try:
            obj, end = decoder.raw_decode(text, pos)
            last_obj = obj
            pos = end
        except json.JSONDecodeError:
            break
    return last_obj

def parse_first_json(path):
    try:
        with open(path, 'r') as f:
            text = f.read()
    except FileNotFoundError:
        return None
    decoder = json.JSONDecoder()
    pos = 0
    while pos < len(text) and text[pos] in ' \t\n\r':
        pos += 1
    if pos >= len(text):
        return None
    try:
        obj, _ = decoder.raw_decode(text, pos)
        return obj
    except json.JSONDecodeError:
        return None

# Parse cell-to-name mapping from scheduler log
cell_name_map = {}
with open('$LOG_FILE', 'r') as f:
    for line in f:
        m = re.search(r'Created cell (\d+) for cgroup \S+/(test_cell_rebal_\w+)', line)
        if m:
            cell_name_map[m.group(1)] = m.group(2)

busy_cell_id = None
idle_cell_id = None
for cid, name in cell_name_map.items():
    if 'busy' in name:
        busy_cell_id = cid
    elif 'idle' in name:
        idle_cell_id = cid

if busy_cell_id is None or idle_cell_id is None:
    print('FAIL:could not find busy/idle cell IDs in scheduler log (found: %s)' % cell_name_map)
    sys.exit(0)

# Use the first monitor snapshot (workload is still hot)
after = parse_first_json('$monitor_output')
baseline = parse_last_json('$baseline_output')

if after is None:
    print('NO_DATA')
    sys.exit(0)

rebalance_count = after.get('rebalance_count', 0)
cells = after.get('cells', {})

if rebalance_count == 0:
    print('FAIL:rebalance_count is 0 (no rebalancing occurred)')
    sys.exit(0)

busy_cpus = cells.get(busy_cell_id, {}).get('num_cpus', 0)
idle_cpus = cells.get(idle_cell_id, {}).get('num_cpus', 0)

if busy_cpus <= idle_cpus:
    print('FAIL:busy cell (%s) has %d CPUs, idle cell (%s) has %d CPUs (expected busy > idle)' % (busy_cell_id, busy_cpus, idle_cell_id, idle_cpus))
    sys.exit(0)

# Check that CPU counts actually changed from baseline (utilization caused re-assignment)
if baseline is not None:
    baseline_cells = baseline.get('cells', {})
    if busy_cell_id in baseline_cells:
        baseline_busy_cpus = baseline_cells[busy_cell_id].get('num_cpus', 0)
        if baseline_busy_cpus > 0 and busy_cpus <= baseline_busy_cpus:
            print('FAIL:busy cell CPUs did not increase from baseline (%d -> %d)' % (baseline_busy_cpus, busy_cpus))
            sys.exit(0)

msg = 'PASS:rebalance_count=%d,busy_cell=%s(%d cpus),idle_cell=%s(%d cpus)' % (
    rebalance_count, busy_cell_id, busy_cpus, idle_cell_id, idle_cpus)
print(msg)
" 2>&1)

    log_info "Rebalancing result: $result"

    case "$result" in
        PASS:*)
            record_result "cpu_rebalancing_demand" "PASSED"
            return 0
            ;;
        FAIL:*)
            local reason="${result#FAIL:}"
            log_error "Rebalancing test failed: $reason"
            record_result "cpu_rebalancing_demand" "FAILED"
            return 1
            ;;
        NO_DATA)
            log_error "Rebalancing test: no monitor data captured"
            record_result "cpu_rebalancing_demand" "FAILED"
            return 1
            ;;
        *)
            log_error "Rebalancing test: unexpected result: $result"
            record_result "cpu_rebalancing_demand" "FAILED"
            return 1
            ;;
    esac
}

# Test: Cpuset change detection - verify scheduler reacts to cpuset.cpus changes
test_cpuset_change() {
    log_info "=== Running test: Cpuset Change Detection ==="

    # Kill existing scheduler if running
    if [[ -n "$SCHED_PID" ]] && ps -p "$SCHED_PID" > /dev/null 2>&1; then
        sudo kill "$SCHED_PID" 2>/dev/null || true
        sleep 2
    fi
    sudo pkill -9 scx_mitosis 2>/dev/null || true
    sleep 2

    # Create two test cells
    local cell_a="$CGROUP_BASE/test_cell_cpuset_a"
    local cell_b="$CGROUP_BASE/test_cell_cpuset_b"

    # Clean up ALL test cgroups (not just cpuset ones) since earlier tests
    # may have left stale cgroups that would confuse the scheduler.
    for cell_name in test_cell_1 test_cell_2 test_cell_3 test_cell_dynamic test_cell_temp test_cell_reuse_a test_cell_reuse_b test_cell_borrow_busy test_cell_borrow_idle test_cell_rebal_busy test_cell_rebal_idle test_cell_cpuset_a test_cell_cpuset_b; do
        local path="$CGROUP_BASE/$cell_name"
        if [[ -d "$path" ]]; then
            for pid in $(cat "$path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done

    sudo mkdir -p "$cell_a"
    sudo mkdir -p "$cell_b"

    # Enable cpuset controller for children of test.slice
    if ! grep -q "cpuset" "$CGROUP_BASE/cgroup.subtree_control" 2>/dev/null; then
        echo "+cpuset" | sudo tee "$CGROUP_BASE/cgroup.subtree_control" > /dev/null
    fi

    # Get total CPU count and split into ranges, leaving some CPUs for cell 0
    local total_cpus=$(nproc)
    local half=$((total_cpus / 2))
    if [[ $total_cpus -lt 4 ]]; then
        log_error "Not enough CPUs for cpuset change test (need at least 4)"
        record_result "cpuset_change" "FAILED"
        return 1
    fi

    # Leave one CPU at each boundary for cell 0 (root cell)
    local first_half="0-$((half - 2))"
    local second_half="${half}-$((total_cpus - 2))"

    # Write initial cpusets: A gets first half, B gets second half
    log_info "Setting initial cpusets: A=$first_half, B=$second_half"
    echo "$first_half" | sudo tee "$cell_a/cpuset.cpus" > /dev/null
    echo "$second_half" | sudo tee "$cell_b/cpuset.cpus" > /dev/null

    # Start scheduler
    start_scheduler

    # Start workloads in both cells
    local workers=$((half > 3 ? 3 : half))
    sudo bash -c "echo \$\$ > $cell_a/cgroup.procs && \
        exec stress-ng --cpu $workers --timeout 60s --quiet" &
    sudo bash -c "echo \$\$ > $cell_b/cgroup.procs && \
        exec stress-ng --cpu $workers --timeout 60s --quiet" &
    sleep 3

    # Sample initial CPU assignments and verify cell A uses first half
    log_info "Verifying initial CPU assignments..."
    local initial_ok=1
    for pid in $(cat "$cell_a/cgroup.procs" 2>/dev/null); do
        if [[ -d "/proc/$pid" ]]; then
            local cpu=$(get_process_cpu "$pid")
            if [[ -n "$cpu" ]] && [[ "$cpu" -ge "$half" ]]; then
                log_error "Cell A PID $pid running on CPU $cpu (expected < $half)"
                initial_ok=0
            fi
        fi
    done

    if [[ $initial_ok -eq 0 ]]; then
        sudo pkill -9 stress-ng 2>/dev/null || true
        record_result "cpuset_change" "FAILED"
        return 1
    fi
    log_info "Initial CPU assignments correct"

    # Now swap the cpusets: A gets second half, B gets first half
    log_info "Swapping cpusets: A=$second_half, B=$first_half"
    echo "$second_half" | sudo tee "$cell_a/cpuset.cpus" > /dev/null
    echo "$first_half" | sudo tee "$cell_b/cpuset.cpus" > /dev/null

    # Wait for scheduler to detect and apply the change
    sleep 5

    # Verify scheduler detected the change
    if ! grep -q "Cpuset change detected" "$LOG_FILE"; then
        log_error "Cpuset change not detected in scheduler logs"
        log_info "Scheduler log tail:"
        tail -20 "$LOG_FILE"
        sudo pkill -9 stress-ng 2>/dev/null || true
        record_result "cpuset_change" "FAILED"
        return 1
    fi

    log_info "Cpuset change detected in scheduler logs"

    # Verify cell A tasks now run on the new cpuset (second half).
    # Only check actively-running worker processes (stress-ng-cpu), not the
    # sleeping master process which may not have been rescheduled yet.
    sleep 2
    local swap_ok=1
    local checked=0
    for pid in $(cat "$cell_a/cgroup.procs" 2>/dev/null); do
        if [[ -d "/proc/$pid" ]]; then
            local comm
            comm=$(cat "/proc/$pid/comm" 2>/dev/null) || continue
            if [[ "$comm" != "stress-ng-cpu" ]]; then
                continue
            fi
            local cpu=$(get_process_cpu "$pid")
            if [[ -n "$cpu" ]]; then
                checked=$((checked + 1))
                if [[ "$cpu" -lt "$half" ]]; then
                    log_warn "After swap: Cell A PID $pid running on CPU $cpu (expected >= $half)"
                    swap_ok=0
                fi
            fi
        fi
    done

    # Kill workloads
    sudo pkill -9 stress-ng 2>/dev/null || true

    if [[ $checked -eq 0 ]]; then
        log_error "No processes found to verify after cpuset swap"
        record_result "cpuset_change" "FAILED"
        return 1
    fi

    if [[ $swap_ok -eq 1 ]]; then
        log_info "After cpuset swap, cell A tasks correctly running on new CPUs"
        record_result "cpuset_change" "PASSED"
        return 0
    else
        log_error "Processes still on old CPUs after cpuset swap"
        record_result "cpuset_change" "FAILED"
        return 1
    fi
}

# Test: New cell gets fair share without disrupting demand-weighted distribution
test_new_cell_gets_fair_share() {
    log_info "=== Running test: New Cell Gets Fair Share ==="

    # Kill existing scheduler if running
    if [[ -n "$SCHED_PID" ]] && ps -p "$SCHED_PID" > /dev/null 2>&1; then
        sudo kill "$SCHED_PID" 2>/dev/null || true
        sleep 2
    fi
    sudo pkill -9 scx_mitosis 2>/dev/null || true
    sleep 2

    # Create two test cells
    local busy_path="$CGROUP_BASE/test_cell_demand_a"
    local idle_path="$CGROUP_BASE/test_cell_demand_b"
    local new_path="$CGROUP_BASE/test_cell_demand_c"

    for path in "$busy_path" "$idle_path" "$new_path"; do
        if [[ -d "$path" ]]; then
            for pid in $(cat "$path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done

    # Clean up other test cgroups that might interfere
    for cell_name in test_cell_1 test_cell_2 test_cell_3 test_cell_dynamic test_cell_temp test_cell_reuse_a test_cell_reuse_b test_cell_borrow_busy test_cell_borrow_idle test_cell_rebal_busy test_cell_rebal_idle test_cell_cpuset_a test_cell_cpuset_b; do
        local path="$CGROUP_BASE/$cell_name"
        if [[ -d "$path" ]]; then
            for pid in $(cat "$path/cgroup.procs" 2>/dev/null); do
                sudo kill -9 "$pid" 2>/dev/null || true
            done
            sudo rmdir "$path" 2>/dev/null || true
        fi
    done

    sudo mkdir -p "$busy_path"
    sudo mkdir -p "$idle_path"

    # Start scheduler with borrowing and rebalancing enabled
    log_info "Starting scx_mitosis with --enable-borrowing --enable-rebalancing..."
    sudo "$SCHEDULER_BIN" \
        --cell-parent-cgroup /test.slice \
        --enable-borrowing \
        --enable-rebalancing \
        --rebalance-cooldown-s 3 \
        --rebalance-threshold 10 \
        > "$LOG_FILE" 2>&1 &
    SCHED_PID=$!

    sleep 3

    # Verify scheduler is running
    if ! ps -p "$SCHED_PID" > /dev/null 2>&1; then
        log_error "Scheduler failed to start. Log output:"
        cat "$LOG_FILE"
        record_result "new_cell_fair_share" "FAILED"
        return 1
    fi

    local state=$(cat /sys/kernel/sched_ext/state 2>/dev/null)
    if [[ "$state" != "enabled" ]]; then
        log_error "sched_ext not enabled. State: $state"
        cat "$LOG_FILE"
        record_result "new_cell_fair_share" "FAILED"
        return 1
    fi

    # Start heavy workload in busy cell (cell A)
    local workers=$(nproc)
    log_info "Starting $workers stress-ng pipe workers in busy cell (A)..."
    sudo bash -c "echo \$\$ > $busy_path/cgroup.procs && \
        exec stress-ng --pipe $workers --timeout 60s --quiet" &

    # Wait for demand stats and rebalancing to trigger
    log_info "Waiting for rebalancing to shift CPUs toward busy cell..."
    sleep 10

    # Capture pre-new-cell monitor snapshot showing rebalanced state
    local pre_snapshot="/tmp/scx_mitosis_monitor_pre_newcell.json"
    log_info "Capturing pre-new-cell monitor snapshot..."
    timeout 4 "$SCHEDULER_BIN" --monitor 1 > "$pre_snapshot" 2>/dev/null || true

    # Now create a third cell while workload is still running
    log_info "Creating new cell (C) while workload is running..."
    sudo mkdir "$new_path"

    # Wait for inotify event and config application (not a full rebalance cycle)
    sleep 5

    # Capture post-new-cell monitor snapshot
    local post_snapshot="/tmp/scx_mitosis_monitor_post_newcell.json"
    log_info "Capturing post-new-cell monitor snapshot..."
    timeout 4 "$SCHEDULER_BIN" --monitor 1 > "$post_snapshot" 2>/dev/null || true

    # Kill workloads
    sudo pkill -9 stress-ng 2>/dev/null || true
    sleep 1

    # Validate
    local result
    result=$(python3 -c "
import json, sys, re

def parse_first_json(path):
    try:
        with open(path, 'r') as f:
            text = f.read()
    except FileNotFoundError:
        return None
    decoder = json.JSONDecoder()
    pos = 0
    while pos < len(text) and text[pos] in ' \t\n\r':
        pos += 1
    if pos >= len(text):
        return None
    try:
        obj, _ = decoder.raw_decode(text, pos)
        return obj
    except json.JSONDecodeError:
        return None

# Parse cell-to-name mapping from scheduler log
cell_name_map = {}
with open('$LOG_FILE', 'r') as f:
    for line in f:
        m = re.search(r'Created cell (\d+) for cgroup \S+/(test_cell_demand_\w+)', line)
        if m:
            cell_name_map[m.group(1)] = m.group(2)

busy_cell_id = None
idle_cell_id = None
new_cell_id = None
for cid, name in cell_name_map.items():
    if name == 'test_cell_demand_a':
        busy_cell_id = cid
    elif name == 'test_cell_demand_b':
        idle_cell_id = cid
    elif name == 'test_cell_demand_c':
        new_cell_id = cid

if busy_cell_id is None or idle_cell_id is None:
    print('FAIL:could not find busy/idle cell IDs in scheduler log (found: %s)' % cell_name_map)
    sys.exit(0)

if new_cell_id is None:
    print('FAIL:new cell (test_cell_demand_c) not found in scheduler log')
    sys.exit(0)

pre = parse_first_json('$pre_snapshot')
post = parse_first_json('$post_snapshot')

if pre is None:
    print('FAIL:no pre-new-cell monitor data')
    sys.exit(0)
if post is None:
    print('FAIL:no post-new-cell monitor data')
    sys.exit(0)

pre_cells = pre.get('cells', {})
post_cells = post.get('cells', {})

# Verify rebalancing occurred before adding new cell
pre_rebalance = pre.get('rebalance_count', 0)
if pre_rebalance == 0:
    print('FAIL:no rebalancing occurred before adding new cell')
    sys.exit(0)

pre_busy_cpus = pre_cells.get(busy_cell_id, {}).get('num_cpus', 0)
if pre_busy_cpus == 0:
    print('FAIL:busy cell has 0 CPUs in pre-snapshot')
    sys.exit(0)

# After new cell: busy cell should retain more than average
total_cpus = sum(c.get('num_cpus', 0) for c in post_cells.values())
num_cells = len(post_cells)
avg_cpus = total_cpus / num_cells if num_cells > 0 else 0

post_busy_cpus = post_cells.get(busy_cell_id, {}).get('num_cpus', 0)
post_new_cpus = post_cells.get(new_cell_id, {}).get('num_cpus', 0)

if post_new_cpus < 1:
    print('FAIL:new cell got %d CPUs (expected >= 1)' % post_new_cpus)
    sys.exit(0)

if post_busy_cpus <= avg_cpus:
    print('FAIL:busy cell has %d CPUs after new cell (avg=%.1f, expected more than avg)' % (post_busy_cpus, avg_cpus))
    sys.exit(0)

msg = 'PASS:pre_busy=%d,post_busy=%d,new=%d,avg=%.1f,total=%d,cells=%d' % (
    pre_busy_cpus, post_busy_cpus, post_new_cpus, avg_cpus, total_cpus, num_cells)
print(msg)
" 2>&1)

    log_info "New cell fair share result: $result"

    case "$result" in
        PASS:*)
            record_result "new_cell_fair_share" "PASSED"
            return 0
            ;;
        FAIL:*)
            local reason="${result#FAIL:}"
            log_error "New cell fair share test failed: $reason"
            record_result "new_cell_fair_share" "FAILED"
            return 1
            ;;
        NO_DATA)
            log_error "New cell fair share test: no monitor data captured"
            record_result "new_cell_fair_share" "FAILED"
            return 1
            ;;
        *)
            log_error "New cell fair share test: unexpected result: $result"
            record_result "new_cell_fair_share" "FAILED"
            return 1
            ;;
    esac
}

# Run dynamic tests
run_dynamic_tests() {
    log_info ""
    log_info "======================================"
    log_info "Running Dynamic Cell Tests"
    log_info "======================================"

    test_dynamic_cell_lifecycle
    test_cpu_reallocation
    test_cell_id_reuse
}

# Print test summary
print_summary() {
    echo ""
    echo "======================================"
    echo "Test Summary"
    echo "======================================"

    local passed=0
    local failed=0
    local inconclusive=0

    for test_name in "${!TEST_RESULTS[@]}"; do
        local result="${TEST_RESULTS[$test_name]}"
        case "$result" in
            PASSED)
                echo -e "  ${GREEN}✓${NC} $test_name"
                ((passed++))
                ;;
            FAILED)
                echo -e "  ${RED}✗${NC} $test_name"
                ((failed++))
                ;;
            INCONCLUSIVE)
                echo -e "  ${YELLOW}?${NC} $test_name"
                ((inconclusive++))
                ;;
        esac
    done

    echo ""
    echo "Results: $passed passed, $failed failed, $inconclusive inconclusive"

    if [[ $failed -gt 0 ]]; then
        return 1
    fi
    return 0
}

# Main test flow
main() {
    echo "======================================"
    echo "scx_mitosis Cell Isolation Test"
    echo "======================================"
    echo "Number of cells: $NUM_CELLS"
    echo "Workers per cell: $WORKERS_PER_CELL"
    echo "Test dynamic: $TEST_DYNAMIC"
    echo "Test borrowing: $TEST_BORROWING"
    echo "Test enqueue borrowing: $TEST_ENQUEUE_BORROWING"
    echo "Test rebalancing: $TEST_REBALANCING"
    echo "Test cpuset change: $TEST_CPUSET_CHANGE"
    echo "Test new cell demand: $TEST_NEW_CELL_DEMAND"
    echo "Test all: $TEST_ALL"
    echo ""

    check_prerequisites
    create_cgroups
    start_scheduler

    # Determine which tests to run
    local run_basic=1
    local run_dynamic=0
    local run_borrowing=0
    local run_enqueue_borrowing=0
    local run_rebalancing=0
    local run_cpuset_change=0
    local run_new_cell_demand=0

    if [[ $TEST_ALL -eq 1 ]]; then
        run_basic=1
        run_dynamic=1
        run_borrowing=1
        run_enqueue_borrowing=1
        run_rebalancing=1
        run_cpuset_change=1
        run_new_cell_demand=1
    elif [[ $TEST_DYNAMIC -eq 1 ]]; then
        run_basic=0
        run_dynamic=1
    elif [[ $TEST_BORROWING -eq 1 ]]; then
        run_basic=0
        run_borrowing=1
    elif [[ $TEST_ENQUEUE_BORROWING -eq 1 ]]; then
        run_basic=0
        run_enqueue_borrowing=1
    elif [[ $TEST_REBALANCING -eq 1 ]]; then
        run_basic=0
        run_rebalancing=1
    elif [[ $TEST_CPUSET_CHANGE -eq 1 ]]; then
        run_basic=0
        run_cpuset_change=1
    elif [[ $TEST_NEW_CELL_DEMAND -eq 1 ]]; then
        run_basic=0
        run_new_cell_demand=1
    fi

    # Run basic isolation test
    if [[ $run_basic -eq 1 ]]; then
        test_basic_isolation
    fi

    # Run dynamic tests
    if [[ $run_dynamic -eq 1 ]]; then
        run_dynamic_tests
    fi

    # Run borrowing test (pipe stressor → select_cpu path)
    if [[ $run_borrowing -eq 1 ]]; then
        test_borrowing pipe cpu_borrowing
    fi

    # Run enqueue borrowing test (cpu stressor → enqueue path)
    if [[ $run_enqueue_borrowing -eq 1 ]]; then
        test_borrowing cpu enqueue_borrowing
    fi

    # Run rebalancing test
    if [[ $run_rebalancing -eq 1 ]]; then
        test_cpu_rebalancing_demand
    fi

    # Run cpuset change test
    if [[ $run_cpuset_change -eq 1 ]]; then
        test_cpuset_change
    fi

    # Run new cell demand test
    if [[ $run_new_cell_demand -eq 1 ]]; then
        test_new_cell_gets_fair_share
    fi

    # Print summary and exit
    if print_summary; then
        log_info "All tests completed successfully"
        exit 0
    else
        log_error "Some tests failed"
        exit 1
    fi
}

main
