#!/bin/bash

set -euo pipefail

# --- UE profile map (index 0..10). 0-index kept for convenience ---
# Initialize with "0" to denote no profile chosen.
declare -a UE_PROFILE_MAP
for i in {0..10}; do UE_PROFILE_MAP[$i]=0; done

# ---------------------------[ Random / diurnal Mode ]----------------------------
# UE scaler mode: 'random' (rapid random targets) or 'diurnal' (legacy cosine)
UE_SCALER_MODE="${UE_SCALER_MODE:-random}"
# Interval (sec) for scaler ticks
UE_SCALER_INTERVAL_SEC="${UE_SCALER_INTERVAL_SEC:-30}"

# -------------------------------[ Config ]------------------------------------

ARTIFACT_DIR="/media/nakaolab-go/2nd/artifacts"
mkdir -p "$ARTIFACT_DIR"

POWER_LOG_FILE="${POWER_LOG_FILE:-$ARTIFACT_DIR/power_usage.csv}"
CPU_PROXY_LOG_FILE="${CPU_PROXY_LOG_FILE:-$ARTIFACT_DIR/cpu_proxy.csv}"
RRC_SINCE_STATE="${RRC_SINCE_STATE:-$ARTIFACT_DIR/.rrc_since}"

# xApp (mac stats) container + path inside container
STATS_CONTAINER="${STATS_CONTAINER:-xapp-chatgpt}"
MAC_STATS_FILE_IN_CONTAINER="${MAC_STATS_FILE_IN_CONTAINER:-/tmp/mac_stats.log}"
PDCP_STATS_FILE_IN_CONTAINER="${PDCP_STATS_FILE_IN_CONTAINER:-/tmp/pdcp_stats.log}"

# Intervals (seconds)
RRC_LOG_INTERVAL_SEC="${RRC_LOG_INTERVAL_SEC:-10}"
POWER_LOG_INTERVAL_SEC="${POWER_LOG_INTERVAL_SEC:-5}"
CPU_PROXY_INTERVAL_SEC="${CPU_PROXY_INTERVAL_SEC:-5}"

# Docker compose helper
DC_BIN=(docker compose)
dc() { command "${DC_BIN[@]}" "$@"; }
dc down -v --remove-orphans

# Service names (override if your compose uses different ones)
AMF_SERVICE="${AMF_SERVICE:-oai-amf}"
UPF_SERVICE="${UPF_SERVICE:-oai-upf}"
GNB_SERVICE="${GNB_SERVICE:-oai-gnb}"
XAPP_SERVICE="${XAPP_SERVICE:-xapp-chatgpt}"

# UE service name pattern (for discovery)
UE_NAME_REGEX="${UE_NAME_REGEX:-^oai-nr-ue}"

# Test controls
TEST_MODE="${TEST_MODE:-0}"  # 1=manual profiles allowed (TEST_* honored); 0=auto
TEST_UES="${TEST_UES:-}"
TEST_PROFILES="${TEST_PROFILES:-}"            # profile name
TEST_DURATION_SEC="${TEST_DURATION_SEC:-0}"  # 0 = run indefinitely

# Diurnal scaling controls (normal mode only; ignored if TEST_UE is set)
UE_MIN="${UE_MIN:-1}"              # Minimum number of UEs to keep connected
UE_MAX="${UE_MAX:-10}"
UE_PEAK_HOUR="${UE_PEAK_HOUR:-17}" # Local hour with peak load (0..23)
UE_NOISE_FRAC="${UE_NOISE_FRAC:-0.15}" # ±fractional noise applied to target

# -------------------------[ Power Model (mW) ]--------------------------------

declare -A PWR_MW_PROFILE=(
    [idle]=60            # screen on, modem camped, no transfers
    [iot_chatty]=150     # tiny bursts, long sleeps
    [browsing]=380       # short transfers with think-time
    [voip]=520           # steady uplink/downlink small packets
    [video]=850          # continuous DL; 720–1080p
    [gaming]=950         # bursty DL + UL keepalives
    [boring]=500         # for single PRB testing
    [bulk_download]=1100 # sustained UL/DL at MCS ceiling
)

# -------------------------------[ Utils ]-------------------------------------

ts_now() { date +%s; }
ts_iso() { date -Is; }
ts_tag() { date +%Y%m%d_%H%M%S; }
log() { echo "[$(ts_iso)] $*"; }

exec_safe() {
    local svc="$1"; shift
    dc exec -T "$svc" bash -lc "$*"
}

# -----------------------------[ Discovery ]-----------------------------------

list_all_services() { dc config --services; }

get_all_ue_services() { list_all_services | grep -E "${UE_NAME_REGEX}" || true; }

get_core_services() {
    # everything except UEs, gNB, xApp
    list_all_services | grep -Ev "(${UE_NAME_REGEX}|^${GNB_SERVICE}$|^${XAPP_SERVICE}$)" || true
}

get_running_ues() {
    dc ps --services --filter "status=running" \
        | grep -E "${UE_NAME_REGEX}" || true
}

# -------------------------[ Readiness helpers ]-------------------------------

svc_cid() { dc ps -q "$1"; }  # container id for a service

wait_healthy() {
    local svc="$1" timeout="${2:-60}"
    local cid; cid="$(svc_cid "$svc")"
    if [[ -z "$cid" ]]; then
        log "wait_healthy: service '$svc' not running."
        return 1
    fi
    local health
    health="$(docker inspect -f '{{if .State.Health}}{{.State.Health.Status}}{{end}}' "$cid" 2>/dev/null || true)"
    if [[ -z "$health" ]]; then
        log "wait_healthy: no healthcheck for '$svc'; skipping."
        return 0
    fi
    local start; start="$(ts_now)"
    while :; do
        health="$(docker inspect -f '{{.State.Health.Status}}' "$cid" 2>/dev/null || echo "")"
        if [[ "$health" == "healthy" ]]; then
            log "$svc is healthy."
            return 0
        fi
        if (( $(ts_now) - start > timeout )); then
            log "wait_healthy timeout for '$svc' (last='$health')."
            return 1
        fi
        sleep 1
    done
}

wait_log() {
    local svc="$1" pat="$2" timeout="${3:-60}"
    local cid; cid="$(svc_cid "$svc")"
    if [[ -z "$cid" ]]; then
        log "wait_log: service '$svc' not running."
        return 1
    fi
    local start; start="$(ts_now)"
    while :; do
        if dc logs "$svc" --since=2m 2>&1 | grep -Eiq -- "$pat"; then
            return 0
        fi
        if (( $(ts_now) - start > timeout )); then
            log "wait_log timeout for '$svc' waiting '$pat'."
            return 1
        fi
        sleep 1
    done
}

get_container_ip() {
    local cid; cid="$(svc_cid "$1")"
    [[ -z "$cid" ]] && return 1
    docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$cid"
}

# --------------------------[ Staged bring-up ]--------------------------------

start_services_list() {
    [[ $# -eq 0 ]] && return 0
    dc up -d "$@" >/dev/null 2>&1
}

start_ue_list() {
    [[ $# -eq 0 ]] && return 0
    for svc in "$@"; do
        log "Starting: $svc"
        dc up -d "$svc" >/dev/null 2>&1
    done
}

start_in_stages() {
    # 1) CORE
    mapfile -t core_svcs < <(get_core_services)
    if [[ ${#core_svcs[@]} -gt 0 ]]; then
        start_services_list "${core_svcs[@]}"
    else
        log "No CORE services to start."
        exit 1
    fi

    # Ensure AMF is ready
    wait_healthy "${AMF_SERVICE}" 15
    wait_healthy "${UPF_SERVICE}" 15
    # exec_safe "${UPF_SERVICE}" 'bash -lc "
    #              set -e                 
    #              sysctl -w net.ipv4.ip_forward=1
    #              sysctl -w net.ipv4.conf.all.rp_filter=2
    #              sysctl -w net.ipv4.conf.eth0.rp_filter=2
    #              sysctl -w net.ipv4.conf.eth1.rp_filter=2
    #              sysctl -w net.ipv4.conf.tun0.rp_filter=2
                 
    #              ip rule add pref 100 iif tun0 to 192.168.71.0/24 lookup 201 || true
    #              ip route replace 192.168.71.0/24 dev eth0 table 201
                 
    #              ip rule add pref 200 iif tun0 lookup 200 || true
    #              ip route replace 12.1.1.0/24 dev tun0 table 200
    #              ip route replace default via 192.168.72.129 dev eth1 table 200
                 
    #              # UE -> DN (192.168.72.128/26) via eth1
    #              iptables -t nat -C POSTROUTING -s 12.1.1.0/24 -d 192.168.72.128/26 -o eth1 -j MASQUERADE 2>/dev/null || \
    #              iptables -t nat -A POSTROUTING -s 12.1.1.0/24 -d 192.168.72.128/26 -o eth1 -j MASQUERADE
                 
    #              # UE -> 192.168.71.0/24 via eth0
    #              iptables -t nat -C POSTROUTING -s 12.1.1.0/24 -d 192.168.71.0/24 -o eth0 -j MASQUERADE 2>/dev/null || \
    #              iptables -t nat -A POSTROUTING -s 12.1.1.0/24 -d 192.168.71.0/24 -o eth0 -j MASQUERADE
    #              "'
    wait_healthy oai-flexric 15
    
    # 2) gNB
    start_services_list "${GNB_SERVICE}"
    wait_log "${GNB_SERVICE}" "Received NGSetupResponse from AMF" 20

    wait_log "${GNB_SERVICE}" "Transaction ID E2 SETUP-REQUEST 0 E2 SETUP-RESPONSE 0" 15

    # 3) xApp
    start_services_list "${XAPP_SERVICE}"
    wait_healthy "${XAPP_SERVICE}" 15

    # 4) UE(s)
    # if [[ -n "${TEST_UE:-}" ]]; then
    #     # Allow multiple UEs in TEST_UE (comma and/or space separated)
    #     IFS=', ' read -r -a _ue_list <<< "${TEST_UE}"
    #     start_services_list "${_ue_list[@]}"
    #     for ue in "${_ue_list[@]}"; do
    #         [[ -z "$ue" ]] && continue
    #         if wait_oai_tun "$ue" 20; then
    #             exec_safe "$ue" 'ip route replace default dev oaitun_ue1' || true
    #             run_traffic_profile "$ue" "random" "0" || true
    #         else
    #             exit 1
    #         fi
    #     done
    # else
    #     # Initial diurnal target, then enforce UE count
    #     local target; if [[ "${UE_SCALER_MODE}" == "diurnal" ]]; then target="$(calc_target_ue_count)"; else target="$(random_target_ue_count)"; fi
    #     log "Initial diurnal target UE count: ${target}"
    #     ensure_ue_count "$target"
    # fi
}

# -------------------------[ Diurnal UE scaling ]------------------------------

calc_target_ue_count() {
    local total_ues="${UE_MAX:-}"
    if [[ -z "$total_ues" ]]; then
        total_ues="$(get_all_ue_services | wc -l | tr -d ' ')"
    fi
    local min="${UE_MIN:-1}"
    if (( total_ues < min )); then total_ues="$min"; fi

    local H M
    H="$(date +%H)"; M="$(date +%M)"

    # Compute target via cosine wave (peak at UE_PEAK_HOUR), then add noise and clamp
    awk -v min="$min" -v max="$total_ues" -v peak="${UE_PEAK_HOUR:-14}" \
        -v noise="${UE_NOISE_FRAC:-0.15}" -v H="$H" -v M="$M" 'BEGIN{
    # time-of-day in days [0,1)
    t = (H*60.0 + M)/1440.0
    ph = peak/24.0
    norm = (cos(2*3.141592653589793*(t - ph)) + 1)/2  # 0..1
    base = min + norm*(max-min)

    srand(); u = (rand()*2.0 - 1.0) * noise  # [-noise, +noise]
    val = base * (1.0 + u)

    if (val < min) val = min
    if (val > max) val = max

    # round to nearest integer
    printf("%.0f\n", val)
  }'
}

ensure_ue_count() {
    local target="$1"

    # Build arrays
    mapfile -t all_ues < <(get_all_ue_services)
    mapfile -t running < <(get_running_ues)

    # Current and desired counts
    local cur="${#running[@]}"

    if (( cur == target )); then
        log "UE scaler: already at target (${target})."
        return 0
    fi

    # Set helpers for membership tests
    declare -A IS_RUNNING=()
    for u in "${running[@]}"; do IS_RUNNING["$u"]=1; done

    # Determine non-running UEs
    declare -a not_running=()
    for u in "${all_ues[@]}"; do
        [[ -z "${IS_RUNNING[$u]:-}" ]] && not_running+=("$u")
    done

    # Start more UEs if needed
    if (( cur < target )); then
        local need=$(( target - cur ))
        log "UE scaler: need +${need} UEs (cur=${cur}, target=${target})."

        if (( ${#not_running[@]} == 0 )); then
            log "UE scaler: no available UEs to start."
            return 0
        fi
        
        local to_start=()
        if command -v shuf >/dev/null 2>&1; then
            mapfile -t to_start < <(printf '%s\n' "${not_running[@]}" | shuf -n "$need" || true)
        else
            to_start=("${not_running[@]:0:$need}")
        fi
        if (( ${#to_start[@]} > 0 )); then
            start_ue_list "${to_start[@]}"
            for ue in "${to_start[@]}"; do
                if wait_log "$ue" "oaitun_ue1 successfully configured" 5 &&
                        exec_safe "$ue" 'ip link show oaitun_ue1 >/dev/null 2>&1'; then
                    exec_safe "$ue" 'ip route replace default dev oaitun_ue1'
                    # run_traffic_profile "$ue" "$(choose_random_profile)" "0"
                    run_traffic_profile "$ue" "boring"
                else
                    dc down -v "$ue" >/dev/null 2>&1 &
                fi
            done            
        fi
        return 0
    fi

    # Otherwise stop some UEs to shrink
    local drop=$(( cur - target ))
    log "UE scaler: need -${drop} UEs (cur=${cur}, target=${target})."
    local to_stop=()
    if command -v shuf >/dev/null 2>&1; then
        mapfile -t to_stop < <(printf '%s\n' "${running[@]}" | shuf -n "$drop" || true)
    else
        to_stop=("${running[@]:0:$drop}")
    fi
    for ue in "${to_stop[@]}"; do
        dc down -v "$ue" >/dev/null 2>&1 &
    done
    wait
}

# --------------------------[ UE traffic helpers ]-----------------------------

PROFILES=(boring browsing video bulk_download voip iot_chatty gaming)
choose_random_profile() {
  if command -v shuf >/dev/null 2>&1; then
    shuf -e "${PROFILES[@]}" -n1
  else
    idx=$(( RANDOM % ${#PROFILES[@]} ))
    echo "${PROFILES[$idx]}"
  fi
}

run_traffic_profile() {
    local ue="$1" profile="$2"
    
    # Record selected profile for this UE into UE_PROFILE_MAP
    if [[ "$ue" =~ ([0-9]+)$ ]]; then
        __idx="${BASH_REMATCH[1]}"
        UE_PROFILE_MAP[$__idx]="$profile"
    fi
    dc cp /home/nakaolab-go/OAI/openairinterface5g/tools/scripts/traffic_runner.sh "$ue":/tmp/traffic_runner.sh >/dev/null 2>&1
    exec_safe "$ue" "bash -lc '
                    prof=\$1;
                    /tmp/traffic_runner.sh --profile \$prof &
                    ' -- $profile"
}

# --------------------------[ RRC / MAC scraping ]-----------------------------

rotate_stats_capture() {
    local tag="${1:-}"
    [[ -z "$tag" ]] && tag="$(ts_tag)"
    if dc ps --services --filter "status=running" | grep -q "^${STATS_CONTAINER}$"; then
        log "Rotating stats from ${STATS_CONTAINER}"
        dc cp "${STATS_CONTAINER}:${MAC_STATS_FILE_IN_CONTAINER}" \
           "${ARTIFACT_DIR}/mac_stats_${tag}.log" || true
        dc cp "${STATS_CONTAINER}:${PDCP_STATS_FILE_IN_CONTAINER}" \
           "${ARTIFACT_DIR}/pdcp_stats_${tag}.log" || true
        # exec_safe "$STATS_CONTAINER" "> '${MAC_STATS_FILE_IN_CONTAINER}'"
        # exec_safe "$STATS_CONTAINER" "> '${PDCP_STATS_FILE_IN_CONTAINER}'"
    else
        log "stats container ${STATS_CONTAINER} not running; skipping rotate"
    fi
}

_since_arg_from_file() {
    local f="$1"
    local val
    val="$(tr -d '[]\r\n' < "$f" 2>/dev/null || true)"
    if [[ -n "$val" ]]; then
        printf -- '--since=%s' "$val"
    else
        printf -- '--since=2m'
    fi
}

dump_gnb_logs() {
    local tag="$1"
    dc logs "${GNB_SERVICE}" 2>&1 | grep -E "NR_RRC" > "${ARTIFACT_DIR}/gnb_rrc_${tag}.log"
    dc logs "${GNB_SERVICE}" 2>&1 | grep -E "\[NR_MAC\] I Frame\.Slot|MAC:" > "${ARTIFACT_DIR}/gnb_${tag}.log"
}

# -----------------------------[ Power / CPU ]---------------------------------

init_power_logs() {
    if [[ ! -s "$POWER_LOG_FILE" ]]; then
        echo "epoch,ue,profile,power_mW,dt_s,energy_mJ" > "$POWER_LOG_FILE"
    fi
    if [[ ! -s "$CPU_PROXY_LOG_FILE" ]]; then
        echo "epoch,container,cpu_percent,mem_usage" > "$CPU_PROXY_LOG_FILE"
    fi
}

profile_total_mw() {
    local profile="$1"
    echo "${PWR_MW_PROFILE[$profile]:-0}"
}

sample_power_once() {
    local epoch="$(ts_now)"
    local dt_s="${POWER_LOG_INTERVAL_SEC}"
    
    while IFS= read -r ue; do
        [[ -z "$ue" ]] && continue
        # Determine profile without container lock by consulting UE_PROFILE_MAP
        local picked_profile picked_duration
        if [[ "$ue" =~ ([0-9]+)$ ]]; then
            __idx="${BASH_REMATCH[1]}"
            picked_profile="${UE_PROFILE_MAP[$__idx]}"
        fi
        picked_duration=""

        if [[ -z "$picked_profile" ]]; then
            if [[ "${TEST_MODE}" -eq 1 ]]; then picked_profile="${TEST_PROFILE}"; else picked_profile="idle"; fi
        fi

        while [[ "$picked_profile" == -* ]]; do
            picked_profile="${picked_profile#-}"
        done
        picked_profile="${picked_profile,,}"

        local p_mw
        p_mw="$(profile_total_mw "$picked_profile")"
        : "${p_mw:=0}"

        local energy_mJ=$(( p_mw * dt_s ))

        echo "${epoch},${ue},${picked_profile},${p_mw},${dt_s},${energy_mJ}" >> "$POWER_LOG_FILE"
    done < <(get_running_ues)
}

# --- CPU proxy: query docker stats by container ID and map back to service name ---
cpu_proxy_once() {
    local epoch="$(ts_now)"

    # Build list of UE service names (same as container names)
    declare -a SVC_LIST=()
    while IFS= read -r svc; do
        [[ -z "$svc" ]] && continue
        SVC_LIST+=("$svc")
    done < <(get_running_ues)

    # Query docker stats by NAME and key 'seen' by NAME
    local fmt='{{.Name}},{{.CPUPerc}},{{.MemUsage}}'
    declare -A SEEN=()

    while IFS=',' read -r name cpu mem; do
        [[ -z "$name" ]] && continue
        name="${name#/}"
        SEEN["$name"]=1
        cpu="${cpu%%%}"  # drop trailing '%'
        echo "${epoch},${name},${cpu},${mem}" >> "$CPU_PROXY_LOG_FILE"
    done < <(docker stats --no-stream --format "$fmt" "${SVC_LIST[@]}" 2>/dev/null || true)

    # Any container absent from docker stats output → record zeros once
    local svc
    for svc in "${SVC_LIST[@]}"; do
        if [[ -z "${SEEN[$svc]:-}" ]]; then
            echo "${epoch},${svc},0,0" >> "$CPU_PROXY_LOG_FILE"
        fi
    done
}

rotate_csv_if_any() {
    local src="$1" prefix="$2" tag="${3:-}"
    [[ -z "$tag" ]] && tag="$(ts_tag)"
    if [[ -s "$src" ]]; then
        mv "$src" "${ARTIFACT_DIR}/${prefix}_${tag}.csv" || true
        case "$prefix" in
            power_usage) echo "epoch,ue,profile,power_mW,dt_s,energy_mJ" > "$src" ;;
            cpu_proxy)   echo "epoch,container,cpu_percent,mem_usage" > "$src" ;;
        esac
    fi
}

finalize_and_rotate_artifacts() {
    local tag; tag="$(ts_tag)"
    rotate_stats_capture "$tag" || true
    dump_gnb_logs "$tag" || true
    rotate_csv_if_any "$POWER_LOG_FILE" "power_usage" "$tag"
    rotate_csv_if_any "$CPU_PROXY_LOG_FILE" "cpu_proxy" "$tag"
}

# -------------------------------[ Signals ]-----------------------------------

__SHUTTING_DOWN=0
on_interrupt() {
    if [[ "${__SHUTTING_DOWN}" -eq 1 ]]; then return; fi
    __SHUTTING_DOWN=1
    log "Signal received (INT/TERM); taking final samples and rotating artifacts…"
    # One last sample tick
    sample_power_once || true
    cpu_proxy_once || true
    # Rotate artifacts and clean up
    finalize_and_rotate_artifacts
    cleanup
    # 130 is the conventional exit code for SIGINT
    exit 130
}

# --------------------------------[ Cleanup ]----------------------------------

cleanup() {
    # Default: stop only UEs; keep infra up for inspection
    log "Cleanup: stopping UE containers only"
    while IFS= read -r ue; do
        [[ -z "$ue" ]] && continue
        dc down -v "$ue" >/dev/null 2>&1 &
    done < <(dc ps --services | grep -E "${UE_NAME_REGEX}" || true)
}

# ---------------------------------[ Main ]------------------------------------

random_target_ue_count() {
    local cur="$(get_running_ues | wc -l | tr -d ' ')"
    bias=$((RANDOM % 100))  # 0–99
    if (( cur <= UE_MIN )); then
        if (( bias < 70 )); then
            delta=$(( (RANDOM % 3) + 1 ))    # +1..+3
        else
            delta=$(( RANDOM % 7 - 3 ))      # -3..+3
        fi
    elif (( cur >= UE_MAX )); then
        if (( bias < 70 )); then
            delta=$(( -((RANDOM % 3) + 1) )) # -1..-3
        else
            delta=$(( RANDOM % 7 - 3 ))      # -3..+3
        fi
    else # normal
        delta=$(( RANDOM % 7 - 3 ))          # -3..+3
    fi
    local target=$((cur + delta))
    if (( target < UE_MIN )); then
        target="${UE_MIN}"
    elif (( target > UE_MAX )); then
        target="${UE_MAX}"
    fi
    
    echo "$target"
}

main() {
    trap on_interrupt INT TERM

    init_power_logs

    # Staged bring-up with readiness checks for both normal and TEST modes
    start_in_stages

    # Kick off traffic workload(s)
    if [[ "${TEST_MODE}" -eq 1 ]]; then
        log "TEST UE(S) ENABLED"
        # Support multiple UEs (comma/space/newline separated)
        read -r -a _ue_list <<< "$(tr ',\n' '  ' <<< "${TEST_UES}")"
        read -r -a _ue_profiles <<< "$(tr ',\n' '  ' <<< "${TEST_PROFILES}")"
        start_ue_list "${_ue_list[@]}"
        for i in "${!_ue_list[@]}"; do
            ue="${_ue_list[$i]}"
            prof="${_ue_profiles[$i]}"
            [[ -z "$ue" ]] && continue
            if wait_log "$ue" "oaitun_ue1 successfully configured" 5 &&
                    exec_safe "$ue" 'ip link show oaitun_ue1 >/dev/null 2>&1'; then
                exec_safe "$ue" "ip route replace default dev oaitun_ue1"
                run_traffic_profile "$ue" "$prof"
            else
                exit 1
            fi
        done
    fi

    # Optional finite duration for TEST mode
    if [[ "${TEST_MODE}" -eq 1 && -n "${TEST_UE:-}" && "${TEST_DURATION_SEC:-0}" -gt 0 ]]; then
        RUN_STOP_TS=$(( $(ts_now) + TEST_DURATION_SEC ))
    fi

    # Main sampling & scaling loop
    local last_rrc="$(( $(ts_now) - RRC_LOG_INTERVAL_SEC ))"
    local last_power="$(( $(ts_now) - POWER_LOG_INTERVAL_SEC ))"
    local last_cpu="$(( $(ts_now) - CPU_PROXY_INTERVAL_SEC ))"
    local last_scale="$(( $(ts_now) - UE_SCALER_INTERVAL_SEC ))"

    log "Entering main loop (sampling + diurnal scaler)…"
    while :; do
        local now; now="$(ts_now)"

        if (( now - last_power >= POWER_LOG_INTERVAL_SEC )); then
            sample_power_once || true
            last_power="$now"
        fi

        if (( now - last_cpu >= CPU_PROXY_INTERVAL_SEC )); then
            cpu_proxy_once || true
            last_cpu="$now"
        fi

        # Re-evaluate diurnal target (normal mode only)
        if [[ "${TEST_MODE}" -eq 0 ]] && (( now - last_scale >= UE_SCALER_INTERVAL_SEC )); then
            local target;
            if [[ "${UE_SCALER_MODE}" == "diurnal" ]]; then
                target="$(calc_target_ue_count)";
            else
                target="$(random_target_ue_count)";
            fi
            log "Scaler tick: target = ${target}"
            ensure_ue_count "$target"
            last_scale="$now"
        fi

        # TEST duration stop gate
        if [[ -n "${RUN_STOP_TS:-}" && "$now" -ge "${RUN_STOP_TS}" ]]; then
            log "TEST run reached duration (${TEST_DURATION_SEC}s); finalizing artifacts…"
            finalize_and_rotate_artifacts
            cleanup
            exit 0
        fi

        sleep 1
    done
}

main "$@"