#!/bin/bash
# Iperf-exclusive traffic generator for OAI UE containers.
# Profiles: browsing, video, bulk_download, voip, iot_chatty, gaming, idle
# Usage:
#   traffic_runner --profile browsing --duration 0 --iperf-server 10.0.0.1:5201
#
set -euo pipefail

# Clear exp.sh lock on exit so sessions can restart later
trap 'rm -f /tmp/traffic_profile.lock' EXIT

DEFAULT_IPERF_PORT="${DEFAULT_IPERF_PORT:-5201}"
TRAFFIC_NET_GATEWAY="${TRAFFIC_NET_GATEWAY:-192.168.71.130}"

if [[ -z "${IPERF_SERVER:-}" ]]; then
    # Candidate hostnames (Docker DNS on the same network) then the gateway IP
    IPERF_HOST_CANDIDATES=${IPERF_HOST_CANDIDATES:-"iperf iperf3 iperf-server ${TRAFFIC_NET_GATEWAY}"}
    for h in $IPERF_HOST_CANDIDATES; do
        if getent hosts "$h" >/dev/null 2>&1; then
            IPERF_SERVER="${h}:${DEFAULT_IPERF_PORT}"
            break
        fi
        # Accept numeric IPv4 even if DNS not present
        if [[ "$h" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
            IPERF_SERVER="${h}:${DEFAULT_IPERF_PORT}"
            break
        fi
    done
fi

# Final check
[[ -n "${IPERF_SERVER:-}" ]] || die "No iperf server found. Set --iperf-server or connect a server to rfsim5g-oai-traffic-net."
IPERF_HOST="${IPERF_SERVER%:*}"
IPERF_PORT="${IPERF_SERVER#*:}"
[[ "$IPERF_PORT" =~ ^[0-9]+$ ]] || die "Invalid iperf port in $IPERF_SERVER"

PROFILE=""
DURATION=0
LOG_FILE="/tmp/traffic_runner.log"

log() { printf '[%(%Y-%m-%dT%H:%M:%S%z)T] %s\n' -1 "$*" | tee -a "$LOG_FILE" >&2; }
die() { log "ERROR: $*"; exit 2; }

# Parse args
while [[ $# -gt 0 ]]; do
    case "$1" in
        --profile) PROFILE="$2"; shift 2;;
        --duration) DURATION="$2"; shift 2;;
        --iperf-server) IPERF_SERVER="$2"; shift 2;;
        --help|-h) cat <<EOF
Usage: $0 --profile <name> [--duration <sec>] [--iperf-server host:port]
Profiles: browsing, video, bulk_download, voip, iot_chatty, gaming, idle
If duration is 0 or omitted, run indefinitely.
EOF
                   exit 0;;
        *) die "Unknown arg: $1";;
    esac
done

[[ -n "$PROFILE" ]] || die "--profile is required"
[[ -n "$IPERF_SERVER" ]] || die "--iperf-server or IPERF_SERVER env is required (e.g., 192.168.0.2:5201)"

IPERF_HOST="${IPERF_SERVER%:*}"
IPERF_PORT="${IPERF_SERVER#*:}"
[[ "$IPERF_PORT" =~ ^[0-9]+$ ]] || die "Invalid iperf port in $IPERF_SERVER"

# Helpers
run_for_duration() {
    # $1: command string to eval in a loop
    # uses global DURATION
    local start now elapsed
    start=$(date +%s)
    while :; do
        now=$(date +%s)
        elapsed=$(( now - start ))
        if (( DURATION > 0 && elapsed >= DURATION )); then break; fi
        eval "$1"
    done
}

# iperf helpers
iperf_tcp_dl() { iperf3 -c "$IPERF_HOST" -p "$IPERF_PORT" -R -t "$1" -J >/dev/null 2>&1 || true; }
iperf_tcp_ul() { iperf3 -c "$IPERF_HOST" -p "$IPERF_PORT"    -t "$1" -J >/dev/null 2>&1 || true; }

iperf_udp_dl() { iperf3 -c "$IPERF_HOST" -p "$IPERF_PORT" -R -u -b "$2" -t "$1" -J >/dev/null 2>&1 || true; }
iperf_udp_ul() { iperf3 -c "$IPERF_HOST" -p "$IPERF_PORT"    -u -b "$2" -t "$1" -J >/dev/null 2>&1 || true; }

# Profiles
do_browsing() {
    log "Browsing (TCP bursts, short think-time) to $IPERF_SERVER"
    run_for_duration '
    dur=$(( 1 + RANDOM % 3 ))     # 1-3s burst
    iperf_tcp_dl "$dur"
    sleep "0.$((100 + RANDOM % 600))"  # 0.1-0.7s think
  '
}

do_video() {
    log "Video (FHD video 5 Mbps sustained) to $IPERF_SERVER"
    run_for_duration '
    iperf3 -c "'"$IPERF_HOST"'" -p "'"$IPERF_PORT"'" -R -u -b 5M -t 30 -J >/dev/null 2>&1
    sleep 1
  '
}

do_bulk_download() {
    # Heavy TCP downlink with multiple streams
    log "Bulk download (TCP DL 8 parallel, long sustained) to $IPERF_SERVER"
    run_for_duration '
    iperf3 -c "'"$IPERF_HOST"'" -p "'"$IPERF_PORT"'" -R -P 8 -t 20 -w 1M -J >/dev/null 2>&1
    sleep 1
  '
}

do_voip() {
    log "VoIP (UDP bi-di ~100–200 kbps) to $IPERF_SERVER"
    run_for_duration '
    iperf_udp_ul 5 120K &
    iperf_udp_dl 5 120K &
    wait
    sleep 0.5
  '
}

do_iot_chatty() {
    log "IoT chatty (UDP UL ~50–80 kbps bursts) to $IPERF_SERVER"
    run_for_duration '
    iperf_udp_ul 2 60K
    sleep $(( 3 + RANDOM % 7 ))
    '
}

do_gaming() {
    log "Gaming (UDP bi-di ~1–3 Mbps bursts) to $IPERF_SERVER"
    run_for_duration '
    rate_ul=$(( 1 + RANDOM % 3 ))M
    rate_dl=$(( 1 + RANDOM % 3 ))M
    iperf_udp_ul 5 "$rate_ul" &
    iperf_udp_dl 5 "$rate_dl" &
    wait
    sleep 0.5
  '
}

do_idle() {
    log "Idle"
    run_for_duration "sleep 2"
}



# ------------- Randomized pattern (bursts of 30s..minutes) ------------------

# Defaults can be overridden via env:
#   MIN_ACTIVE_SEC / MAX_ACTIVE_SEC — each session duration (default 30..180)
#   MIN_IDLE_SEC / MAX_IDLE_SEC     — wait between sessions (default 5..60)
#   BURST_FILE                      — path watched for system-wide bursts
#   BURST_MAX_AGE_SEC              — consider burst file "fresh" within N sec
MIN_ACTIVE_SEC="${MIN_ACTIVE_SEC:-30}"
MAX_ACTIVE_SEC="${MAX_ACTIVE_SEC:-180}"
MIN_IDLE_SEC="${MIN_IDLE_SEC:-5}"
MAX_IDLE_SEC="${MAX_IDLE_SEC:-60}"
BURST_FILE="${BURST_FILE:-/tmp/system_burst}"
BURST_MAX_AGE_SEC="${BURST_MAX_AGE_SEC:-45}"

_pick_random_profile() {
    # Weighted pick to feel realistic; adjust as desired
    local r=$((RANDOM % 100))
    if   (( r < 40 )); then echo browsing     # 40%
    elif (( r < 55 )); then echo voip         # 15%
    elif (( r < 70 )); then echo video        # 15%
    elif (( r < 85 )); then echo gaming       # 15%
    else                      echo bulk_download # 15%
    fi
}

_run_profile_for() {
    local name="$1" seconds="$2"
    DURATION="$seconds"
    case "$name" in
        browsing)      do_browsing;;
        video)         do_video;;
        bulk_download) do_bulk_download;;
        voip)          do_voip;;
        iot_chatty)    do_iot_chatty;;
        gaming)        do_gaming;;
        idle)          do_idle;;
        *)             do_browsing;;
    esac
}

_is_burst_now() {
    # return 0 if a fresh burst marker exists
    [[ -f "$BURST_FILE" ]] || return 1
    local now mtime age
    now=$(date +%s)
    mtime=$(stat -c %Y "$BURST_FILE" 2>/dev/null || echo 0)
    age=$(( now - mtime ))
    (( age >= 0 && age <= BURST_MAX_AGE_SEC ))
}

do_random() {
    log "Random pattern: independent per-UE bursts to $IPERF_SERVER (session: ${MIN_ACTIVE_SEC}-${MAX_ACTIVE_SEC}s; idle: ${MIN_IDLE_SEC}-${MAX_IDLE_SEC}s)"
    while :; do
        local prof dur idle
        prof="$(_pick_random_profile)"
        dur=$(( MIN_ACTIVE_SEC + RANDOM % (MAX_ACTIVE_SEC - MIN_ACTIVE_SEC + 1) ))

        if _is_burst_now; then
            # If a system burst is signaled, start immediately and slightly extend session
            idle=0
            dur=$(( dur + 20 + RANDOM % 40 ))  # +20..60s
        else
            idle=$(( MIN_IDLE_SEC + RANDOM % (MAX_IDLE_SEC - MIN_IDLE_SEC + 1) ))
        fi

        log "Random session: profile=${prof} duration=${dur}s (then idle ${idle}s)"
        _run_profile_for "$prof" "$dur" || true
        sleep "$idle"
    done
}

case "$PROFILE" in
    browsing)       do_browsing;;
    video)          do_video;;
    bulk_download)  do_bulk_download;;
    voip)           do_voip;;
    iot_chatty)     do_iot_chatty;;
    gaming)         do_gaming;;
    idle)           do_idle;;
    random)        do_random;;
    *) die "Unknown profile: $PROFILE";;
esac