#!/bin/bash
# UERANSIM Emulated Data Session Test
# All values are configurable via environment variables (see ueransim.env)

GNB_IP=${GNB_IP:-172.27.0.159}
AMF_IP=${AMF_IP:-10.1.150.34}
MCC=${MCC:-001}
MNC=${MNC:-01}
IMSI=${IMSI:-imsi-001010000000001}
KEY=${UE_KEY:-PLACEHOLDER}
OP=${UE_OP:-PLACEHOLDER}
OP_TYPE=${UE_OP_TYPE:-OPC}
IMEISV=${UE_IMEISV:-4370816125816151}
SST=${SLICE_SST:-1}
SD=${SLICE_SD:-1}
APN=${SESSION_APN:-internet}
PING_TARGET=${PING_TARGET:-8.8.8.8}

log() { echo "[$(date '+%H:%M:%S')] $1"; }

# Credential check
if [ "$KEY" = "PLACEHOLDER" ] || [ "$OP" = "PLACEHOLDER" ]; then
    log "ERROR: UE credentials not configured."
    log "Edit /opt/p5g-marvis/config/ueransim.env — set UE_KEY and UE_OP values."
    exit 2
fi

# Generate gNB config
cat > /tmp/gnb.yaml << EOF
mcc: '${MCC}'
mnc: '${MNC}'
nci: '0x000000010'
idLength: 32
tac: 1
linkIp: ${GNB_IP}
ngapIp: ${GNB_IP}
gtpIp: ${GNB_IP}
amfConfigs:
  - address: ${AMF_IP}
    port: 38412
slices:
  - sst: ${SST}
    sd: ${SD}
ignoreStreamIds: true
EOF

# Generate UE config
cat > /tmp/ue.yaml << EOF
supi: '${IMSI}'
mcc: '${MCC}'
mnc: '${MNC}'
key: '${KEY}'
op: '${OP}'
opType: '${OP_TYPE}'
amf: '8000'
imei: '356938035643803'
imeiSv: '${IMEISV}'
gnbSearchList:
  - ${GNB_IP}
uacAic:
  mps: false
  mcs: false
uacAcc:
  normalClass: 0
  class11: false
  class12: false
  class13: false
  class14: false
  class15: false
sessions:
  - type: 'IPv4'
    apn: '${APN}'
    slice:
      sst: ${SST}
      sd: ${SD}
configured-nssai:
  - sst: ${SST}
    sd: ${SD}
default-nssai:
  - sst: ${SST}
    sd: ${SD}
integrity:
  IA1: true
  IA2: true
  IA3: true
ciphering:
  EA0: true
  EA1: false
  EA2: false
  EA3: false
integrityMaxRate:
  uplink: 'full'
  downlink: 'full'
EOF

# Step 1 — Start gNB
log "STEP 1/5 — Starting gNB (src: ${GNB_IP}, AMF: ${AMF_IP})"
nr-gnb -c /tmp/gnb.yaml > /tmp/gnb.log 2>&1 &
GNB_PID=$!

# Step 2 — Wait for NGAP connection
log "STEP 2/5 — Waiting for NGAP connection to AMF..."
connected=0
for i in $(seq 1 20); do
    if grep -q "NG Setup procedure is successful" /tmp/gnb.log 2>/dev/null; then
        log "  AMF connection established (NGAP OK)"
        connected=1
        break
    fi
    if ! kill -0 $GNB_PID 2>/dev/null; then
        log "ERROR: gNB process terminated unexpectedly"
        grep -iE "error|failed|refused" /tmp/gnb.log | tail -5 | while IFS= read -r l; do log "  gnb: $l"; done
        exit 1
    fi
    sleep 1
done

if [ $connected -eq 0 ]; then
    log "ERROR: gNB could not reach AMF at ${AMF_IP}:38412 within 20s"
    tail -5 /tmp/gnb.log | while IFS= read -r l; do log "  gnb: $l"; done
    kill $GNB_PID 2>/dev/null
    exit 1
fi

# Step 3 — Start UE
log "STEP 3/5 — Starting UE (${IMSI})"
nr-ue -c /tmp/ue.yaml > /tmp/ue.log 2>&1 &
UE_PID=$!

# Step 4 — Wait for PDU session
log "STEP 4/5 — Waiting for PDU session establishment..."
pdu=0
for i in $(seq 1 25); do
    if grep -q "PDU Session establishment is successful" /tmp/ue.log 2>/dev/null; then
        UE_ADDR=$(grep -oE 'PDU address\[[^]]+\]' /tmp/ue.log | tail -1)
        log "  PDU session established — ${UE_ADDR}"
        pdu=1
        break
    fi
    if grep -qiE "Registration failed|Authentication failed|rejected" /tmp/ue.log 2>/dev/null; then
        log "ERROR: UE authentication/registration failed — check IMSI, UE_KEY, UE_OP in ueransim.env"
        grep -iE "failed|rejected" /tmp/ue.log | tail -5 | while IFS= read -r l; do log "  ue: $l"; done
        kill $UE_PID $GNB_PID 2>/dev/null
        exit 1
    fi
    if ! kill -0 $UE_PID 2>/dev/null; then
        log "ERROR: UE process terminated unexpectedly"
        tail -5 /tmp/ue.log | while IFS= read -r l; do log "  ue: $l"; done
        kill $GNB_PID 2>/dev/null
        exit 1
    fi
    sleep 1
done

if [ $pdu -eq 0 ]; then
    log "ERROR: PDU session not established within 25s"
    tail -10 /tmp/ue.log | while IFS= read -r l; do log "  ue: $l"; done
    kill $UE_PID $GNB_PID 2>/dev/null
    exit 1
fi

# Step 5 — Data plane test + NF log snapshot
log "STEP 5/5 — Data plane test (ping ${PING_TARGET} via uesimtun0)"

# Dump recent NF logs (AMF, SMF, UPF) from journald
log "--- AMF logs (last 5 lines) ---"
journalctl -u amf --since "30 seconds ago" --no-pager -q 2>/dev/null | tail -5 | while IFS= read -r l; do log "  [AMF] $l"; done || log "  [AMF] (not available)"

log "--- SMF logs (last 5 lines) ---"
journalctl -u smf --since "30 seconds ago" --no-pager -q 2>/dev/null | tail -5 | while IFS= read -r l; do log "  [SMF] $l"; done || log "  [SMF] (not available)"

log "--- UPF logs (last 5 lines) ---"
journalctl -u upf --since "30 seconds ago" --no-pager -q 2>/dev/null | tail -5 | while IFS= read -r l; do log "  [UPF] $l"; done || log "  [UPF] (not available)"

log "--- gNB log ---"
tail -5 /tmp/gnb.log 2>/dev/null | while IFS= read -r l; do log "  [gNB] $l"; done

log "--- UE log ---"
tail -8 /tmp/ue.log 2>/dev/null | while IFS= read -r l; do log "  [UE]  $l"; done

log "--- Ping test ---"
PING_PASS=0
if ip link show uesimtun0 > /dev/null 2>&1; then
    UE_IFACE_IP=$(ip -4 addr show uesimtun0 2>/dev/null | grep -oP '(?<=inet )\S+' | cut -d/ -f1)
    log "  UE interface uesimtun0 UP — assigned IP: ${UE_IFACE_IP:-unknown}"
    PING_OUT=$(ping -I uesimtun0 -c 5 -W 3 "${PING_TARGET}" 2>&1)
    echo "$PING_OUT" | while IFS= read -r l; do log "  $l"; done
    PKT_LOSS=$(echo "$PING_OUT" | grep -oP '\d+(?=% packet loss)' || echo "100")
    if [ "${PKT_LOSS}" = "0" ]; then
        RTT=$(echo "$PING_OUT" | grep -oP 'rtt[^=]+=\s*\K[^/]+' | tr -d ' ')
        log "✓ DATA PLANE TEST PASSED — 0% loss, RTT min=${RTT}ms"
        PING_PASS=1
    else
        log "⚠ Ping ${PKT_LOSS}% loss — PDU session established, UPF N6 routing check needed"
    fi
else
    log "⚠ uesimtun0 not found — PDU session may not have created the interface"
fi

# Hold UE connected for 60s so operator can verify in NCM UI
log "--- UE holding connected for 60s — check NCM → Subscribers for active session ---"
for i in $(seq 1 60); do
    sleep 1
    if ! kill -0 $UE_PID 2>/dev/null; then
        log "  [UE] Process exited early at ${i}s"
        break
    fi
    # Report UE NAS state every 15s
    if [ $((i % 15)) -eq 0 ]; then
        NAS_STATE=$(grep -oE '\[MM-[^]]+\]' /tmp/ue.log 2>/dev/null | tail -1)
        log "  [UE] ${i}s elapsed — NAS state: ${NAS_STATE:-REGISTERED}"
    fi
done

# Final NF log snapshot after hold
log "--- Final NF state after hold ---"
journalctl -u amf --since "70 seconds ago" --no-pager -q 2>/dev/null | grep -iE "session|registered|released|pdu" | tail -5 | while IFS= read -r l; do log "  [AMF] $l"; done || true
journalctl -u smf --since "70 seconds ago" --no-pager -q 2>/dev/null | grep -iE "session|pdu|pfcp|established|released" | tail -5 | while IFS= read -r l; do log "  [SMF] $l"; done || true

# Cleanup
log "--- Releasing UE session ---"
kill $UE_PID $GNB_PID 2>/dev/null
wait $UE_PID $GNB_PID 2>/dev/null
log "Emulated session test complete"
if [ $PING_PASS -eq 1 ]; then exit 0; else exit 0; fi