#!/bin/bash
# run_instruments.sh

set -uo pipefail

BINARY="broadcast_server"
FEATURES="b_server"
LOAD_TESTER_BIN="chat_load_tester"
LOAD_TESTER_FEATURES="load"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="instruments_output_${TIMESTAMP}"
COOLDOWN_SECS=20
ONLY_TEMPLATE=""
CAPTURE_DURATION=90
SERVER_READY_TIMEOUT=30

while [[ $# -gt 0 ]]; do
    case "$1" in
        --time-limit) CAPTURE_DURATION="$2"; shift 2 ;;
        --template)   ONLY_TEMPLATE="$2";    shift 2 ;;
        *)            echo "Unknown: $1"; exit 1 ;;
    esac
done

ALL_TEMPLATES=(
    "Time Profiler"
    "CPU Profiler"
    "Counters"
    "System Trace"
    "Thread States"
    "Locks"
    "Swift Concurrency"
    "Allocations"
    "Leaks"
    "VM Tracker"
    "Virtual Memory Trace"
    "Network"
    "File Activity"
    "Disk Monitor"
    "I/O"
    "System Usage"
    "Energy Log"
    "Thermal State"
)

if [[ -n "$ONLY_TEMPLATE" ]]; then
    TEMPLATES=("$ONLY_TEMPLATE")
else
    TEMPLATES=("${ALL_TEMPLATES[@]}")
fi

TEMPLATE_COUNT=${#TEMPLATES[@]}
TOTAL_ESTIMATE=$(( TEMPLATE_COUNT * (CAPTURE_DURATION + COOLDOWN_SECS) ))
TOTAL_MINUTES=$(( TOTAL_ESTIMATE / 60 ))

echo ""
echo "═══════════════════════════════════════════════════════════════"
echo "📡 INSTRUMENTS PROFILING"
echo "   Server:     $BINARY ($FEATURES)"
echo "   Tester:     $LOAD_TESTER_BIN ($LOAD_TESTER_FEATURES)"
echo "   Templates:  $TEMPLATE_COUNT"
echo "   Capture:    ${CAPTURE_DURATION}s per template"
echo "   Cooldown:   ${COOLDOWN_SECS}s between templates"
echo "   Estimated:  ~${TOTAL_MINUTES}m total"
echo "   Traces:     target/instruments/"
echo "═══════════════════════════════════════════════════════════════"
echo ""

if [[ "$(uname)" != "Darwin" ]]; then
    echo "❌ Instruments requires macOS."
    exit 1
fi

if ! cargo instruments --version &>/dev/null 2>&1; then
    echo "❌ cargo-instruments not found."
    echo "   Install: cargo install cargo-instruments"
    exit 1
fi

# ── Kill any existing server on port 8080 ─────────────────────────────────

echo "🔍 Checking port 8080…"
EXISTING_PID=$(lsof -ti tcp:8080 2>/dev/null || true)
if [[ -n "$EXISTING_PID" ]]; then
    echo "   Killing existing process (PID: $EXISTING_PID)…"
    kill "$EXISTING_PID" 2>/dev/null || true
    sleep 3
    STILL=$(lsof -ti tcp:8080 2>/dev/null || true)
    if [[ -n "$STILL" ]]; then
        kill -9 "$STILL" 2>/dev/null || true
        sleep 2
    fi
fi
echo "   ✅ Port 8080 free."
echo ""

# ── Build ─────────────────────────────────────────────────────────────────

echo "🔨 Building $BINARY…"
if ! cargo build --release --features "$FEATURES" --bin "$BINARY" 2>&1 | tail -3; then
    echo "❌ Build failed."
    exit 1
fi
echo "   ✅ Build complete."
echo ""

# ── Setup output ──────────────────────────────────────────────────────────

mkdir -p "$OUTPUT_DIR"
SUMMARY_FILE="${OUTPUT_DIR}/00_summary.txt"

TOTAL=0; PASSED=0; SKIPPED=0; FAILED=0
RESULT_NAMES=()
RESULT_STATUSES=()
RESULT_ELAPSED=()

{
    echo "INSTRUMENTS SESSION"
    echo "Started:   $(date)"
    echo "Binary:    $BINARY"
    echo "Tester:    $LOAD_TESTER_BIN"
    echo "Templates: $TEMPLATE_COUNT"
    echo "Capture:   ${CAPTURE_DURATION}s"
    echo "Traces:    target/instruments/"
    echo ""
} > "$SUMMARY_FILE"

echo "📁 Output:  $OUTPUT_DIR"
echo "📂 Traces:  target/instruments/"
echo ""

# ── Load tester launcher ──────────────────────────────────────────────────

start_load_tester_when_ready() {
    local ready_timeout="$1"
    local log_file="$2"
    (
        local waited=0
        while [[ $waited -lt $ready_timeout ]]; do
            if curl -sf --max-time 2 "http://localhost:8080/stats" >/dev/null 2>&1; then
                echo "[load-tester] Ready after ${waited}s. Starting." >> "$log_file"
                cargo run \
                    --release \
                    --bin "$LOAD_TESTER_BIN" \
                    --features "$LOAD_TESTER_FEATURES" \
                    >> "$log_file" 2>&1
                echo "[load-tester] Done." >> "$log_file"
                return 0
            fi
            sleep 1
            waited=$((waited + 1))
        done
        echo "[load-tester] Server not ready in ${ready_timeout}s." >> "$log_file"
    ) &
    echo $!
}

# ── Per-template loop ─────────────────────────────────────────────────────

for TEMPLATE in "${TEMPLATES[@]}"; do
    TOTAL=$((TOTAL + 1))
    SAFE_NAME=$(echo "$TEMPLATE" | tr ' ' '_')
    LOG_FILE="${OUTPUT_DIR}/${SAFE_NAME}.txt"
    LOAD_LOG="${OUTPUT_DIR}/${SAFE_NAME}_load_tester.txt"

    REMAINING=$(( TEMPLATE_COUNT - TOTAL + 1 ))
    REMAINING_MIN=$(( REMAINING * (CAPTURE_DURATION + COOLDOWN_SECS) / 60 ))

    echo "═══════════════════════════════════════════════════════════════"
    echo "▶  [${TOTAL}/${TEMPLATE_COUNT}] $TEMPLATE"
    echo "   Capture: ${CAPTURE_DURATION}s  |  ~${REMAINING_MIN}m remaining"
    echo "═══════════════════════════════════════════════════════════════"

    {
        echo "PROFILE:   $TEMPLATE"
        echo "Started:   $(date)"
        echo "Capture:   ${CAPTURE_DURATION}s"
        echo ""
    } > "$LOG_FILE"

    # Free port 8080 if stale process exists
    STALE=$(lsof -ti tcp:8080 2>/dev/null || true)
    if [[ -n "$STALE" ]]; then
        echo "   Clearing stale process (PID $STALE)…"
        kill "$STALE" 2>/dev/null || true
        sleep 3
        STILL=$(lsof -ti tcp:8080 2>/dev/null || true)
        [[ -n "$STILL" ]] && { kill -9 "$STILL" 2>/dev/null || true; sleep 2; }
    fi

    # Note the most recent trace before this run so we can find the new one
    LATEST_BEFORE=$(ls -t target/instruments/*.trace 2>/dev/null | head -1 || true)

    # Start load tester in background
    echo "   Starting load tester (waiting for server)…"
    LOAD_TESTER_PID=$(start_load_tester_when_ready "$SERVER_READY_TIMEOUT" "$LOAD_LOG")

    # Run cargo instruments
    # NOTE: no set -e here — non-zero exit for unavailable templates is normal
    START_TIME=$(date +%s)
    STDERR_TMP=$(mktemp)

    cargo instruments \
        -t "$TEMPLATE" \
        --time-limit "$CAPTURE_DURATION" \
        --release \
        --features "$FEATURES" \
        --bin "$BINARY" \
        2>"$STDERR_TMP" | tee -a "$LOG_FILE"
    EXIT_CODE=$?

    STDERR_CONTENT=$(cat "$STDERR_TMP")
    rm -f "$STDERR_TMP"

    END_TIME=$(date +%s)
    ELAPSED=$(( END_TIME - START_TIME ))

    # Kill load tester
    if kill -0 "$LOAD_TESTER_PID" 2>/dev/null; then
        kill "$LOAD_TESTER_PID" 2>/dev/null || true
        wait "$LOAD_TESTER_PID" 2>/dev/null || true
    fi
    pkill -f "$LOAD_TESTER_BIN" 2>/dev/null || true
    sleep 1

    # ── Copy the new trace to output directory ────────────────────────────
    #
    # cargo instruments writes to target/instruments/ with a timestamped name.
    # Find the newest trace that appeared after LATEST_BEFORE.
    TRACE_DEST="${OUTPUT_DIR}/${SAFE_NAME}.trace"
    TRACE_COPIED=false

    if [[ -d "target/instruments" ]]; then
        # Get the newest trace in target/instruments
        NEWEST_TRACE=$(ls -t target/instruments/*.trace 2>/dev/null | head -1 || true)

        if [[ -n "$NEWEST_TRACE" && "$NEWEST_TRACE" != "$LATEST_BEFORE" ]]; then
            echo "   📂 Copying trace: $(basename "$NEWEST_TRACE")"
            cp -r "$NEWEST_TRACE" "$TRACE_DEST"
            TRACE_COPIED=true
            echo "   ✅ Saved as: $TRACE_DEST"
        elif [[ -n "$NEWEST_TRACE" && $EXIT_CODE -eq 0 ]]; then
            # Same name as before — instruments may have reused the directory
            # Copy it anyway if the exit was successful
            echo "   📂 Copying trace (same name as previous): $(basename "$NEWEST_TRACE")"
            cp -r "$NEWEST_TRACE" "$TRACE_DEST"
            TRACE_COPIED=true
        fi
    fi

    if [[ "$TRACE_COPIED" != "true" && $EXIT_CODE -eq 0 ]]; then
        echo "   ⚠️  Could not identify trace file for this run."
    fi

    # Log stderr if present
    if [[ -n "$STDERR_CONTENT" ]]; then
        {
            echo ""
            echo "── stderr ──────────────────────────────────────────────────────"
            echo "$STDERR_CONTENT"
        } >> "$LOG_FILE"
    fi

    # Classify result
    STATUS="FAIL"
    ICON="❌"
    if [[ $EXIT_CODE -eq 0 ]]; then
        STATUS="PASS"; ICON="✅"; PASSED=$((PASSED + 1))
    elif echo "$STDERR_CONTENT" | grep -qi \
        "not found\|not available\|unsupported\|no template\|unknown instrument\|unrecognized"; then
        STATUS="SKIP"; ICON="⚠️ "; SKIPPED=$((SKIPPED + 1))
    else
        STATUS="FAIL"; ICON="❌"; FAILED=$((FAILED + 1))
        echo "   stderr:"
        echo "$STDERR_CONTENT" | head -5 | sed 's/^/      /'
    fi

    {
        echo ""
        echo "Finished:  $(date)"
        echo "Elapsed:   ${ELAPSED}s"
        echo "Exit code: $EXIT_CODE"
        echo "Status:    $STATUS"
        echo "Trace:     $TRACE_COPIED"
    } >> "$LOG_FILE"

    {
        echo "[$TEMPLATE]"
        echo "  Status:    $STATUS"
        echo "  Elapsed:   ${ELAPSED}s"
        echo "  Exit code: $EXIT_CODE"
        echo "  Trace:     $TRACE_COPIED"
        echo ""
    } >> "$SUMMARY_FILE"

    RESULT_NAMES+=("$TEMPLATE")
    RESULT_STATUSES+=("$STATUS")
    RESULT_ELAPSED+=("${ELAPSED}s")

    echo "${ICON} ${STATUS}: $TEMPLATE (${ELAPSED}s)"

    if [[ $TOTAL -lt $TEMPLATE_COUNT ]]; then
        echo "   ⏳ Cooldown ${COOLDOWN_SECS}s…"
        sleep "$COOLDOWN_SECS"
    fi
    echo ""
done

# ── Final summary ─────────────────────────────────────────────────────────

echo "═══════════════════════════════════════════════════════════════"
echo "📊 SESSION COMPLETE — $(date)"
echo "═══════════════════════════════════════════════════════════════"
echo ""

printf "%-30s %-8s %s\n" "Template" "Status" "Elapsed"
printf "%-30s %-8s %s\n" "──────────────────────────────" "────────" "───────"

i=0
while [[ $i -lt ${#RESULT_NAMES[@]} ]]; do
    tpl="${RESULT_NAMES[$i]}"
    status="${RESULT_STATUSES[$i]}"
    elapsed="${RESULT_ELAPSED[$i]}"
    icon="✅"
    [[ "$status" == "SKIP" ]] && icon="⚠️ "
    [[ "$status" == "FAIL" ]] && icon="❌"
    printf "%-30s %s %-4s  %s\n" "$tpl" "$icon" "$status" "$elapsed"
    i=$((i + 1))
done

echo ""
echo "  Total:  $TEMPLATE_COUNT  ✅ Pass: $PASSED  ⚠️  Skip: $SKIPPED  ❌ Fail: $FAILED"
echo ""

{
    echo "SESSION COMPLETE: $(date)"
    echo "Pass=$PASSED Skip=$SKIPPED Fail=$FAILED"
} >> "$SUMMARY_FILE"

echo "📁 Output:  $OUTPUT_DIR"
echo "📂 Traces:  target/instruments/"
echo ""
echo "Open copied traces:"
echo "  open ${OUTPUT_DIR}/*.trace"
echo ""
echo "Or open all traces from target directly:"
echo "  open target/instruments/broadcast_server_*.trace"
echo ""
echo "Re-run one template:"
echo "  ./run_instruments.sh --template \"Time Profiler\" --time-limit 120"

[[ $FAILED -gt 0 ]] && exit 1
exit 0