#!/usr/bin/env bash
set -euo pipefail

SYMCLAW="${SYMCLAW_BIN:-./target/release/symclaw}"
TOTAL="${STRESS_COUNT:-1000}"
CONCURRENCY="${STRESS_CONCURRENCY:-20}"
RESULTS_DIR=$(mktemp -d)
FAILURES=0
OOM_COUNT=0

if [[ ! -x "$SYMCLAW" ]]; then
    echo "ERROR: symclaw binary not found at $SYMCLAW"
    echo "Build with: cargo build --release -p symclaw-cli"
    exit 1
fi

echo "=== SymClaw Stress Test ==="
echo "Binary: $SYMCLAW"
echo "Expressions: $TOTAL (concurrency: $CONCURRENCY)"
echo "Results dir: $RESULTS_DIR"
echo ""

# Expression generators
expressions=(
    "diff(x^2 + 3*x + 1, x)"
    "diff(sin(x)*cos(x), x)"
    "diff(e^(x^2), x)"
    "simplify((x^2 - 1)/(x - 1))"
    "simplify(sin(x)^2 + cos(x)^2)"
    "expand((x + y)^4)"
    "expand((a + b + c)^3)"
    "factor(x^4 - 1)"
    "diff(ln(x^2 + 1), x)"
    "simplify((x^3 - 8)/(x - 2))"
    "diff(tan(x), x)"
    "expand((x + 1)^6)"
    "simplify(x^2*x^3)"
    "diff(x^x, x)"
    "simplify((a^2 - b^2)/(a + b))"
    "expand((1 + x + x^2)^2)"
    "diff(sin(cos(tan(x))), x)"
    "simplify(sqrt(x^2))"
    "diff(e^(sin(x)), x)"
    "expand((x - y)^5)"
)

num_exprs=${#expressions[@]}

run_one() {
    local id=$1
    local expr=${expressions[$((RANDOM % num_exprs))]}
    local outfile="$RESULTS_DIR/$id"

    local start_ns
    start_ns=$(date +%s%N)

    local exit_code=0
    timeout 30 "$SYMCLAW" eval "$expr" > "$outfile.out" 2> "$outfile.err" || exit_code=$?

    local end_ns
    end_ns=$(date +%s%N)
    local elapsed_ms=$(( (end_ns - start_ns) / 1000000 ))

    echo "$id $exit_code $elapsed_ms" >> "$RESULTS_DIR/timings.log"

    if [[ $exit_code -ne 0 ]]; then
        if grep -qi "out of memory\|oom\|alloc" "$outfile.err" 2>/dev/null; then
            echo "OOM" >> "$RESULTS_DIR/oom.log"
        fi
        echo "FAIL" >> "$RESULTS_DIR/failures.log"
    fi
}

export -f run_one
export SYMCLAW RESULTS_DIR expressions num_exprs

# Initialize log files
> "$RESULTS_DIR/timings.log"
> "$RESULTS_DIR/failures.log"
> "$RESULTS_DIR/oom.log"

echo "Running $TOTAL expressions..."
seq 1 "$TOTAL" | xargs -P "$CONCURRENCY" -I{} bash -c 'run_one {}'

# Collect results
FAILURES=$(wc -l < "$RESULTS_DIR/failures.log")
OOM_COUNT=$(wc -l < "$RESULTS_DIR/oom.log")
SUCCESSES=$((TOTAL - FAILURES))

# Calculate latency percentiles
if [[ -s "$RESULTS_DIR/timings.log" ]]; then
    latencies=$(awk '{print $3}' "$RESULTS_DIR/timings.log" | sort -n)
    count=$(echo "$latencies" | wc -l)

    p50_idx=$(( count * 50 / 100 ))
    p95_idx=$(( count * 95 / 100 ))
    p99_idx=$(( count * 99 / 100 ))
    [ "$p50_idx" -lt 1 ] && p50_idx=1
    [ "$p95_idx" -lt 1 ] && p95_idx=1
    [ "$p99_idx" -lt 1 ] && p99_idx=1

    p50=$(echo "$latencies" | sed -n "${p50_idx}p")
    p95=$(echo "$latencies" | sed -n "${p95_idx}p")
    p99=$(echo "$latencies" | sed -n "${p99_idx}p")
    avg=$(echo "$latencies" | awk '{s+=$1}END{printf "%.0f", s/NR}')
    max_lat=$(echo "$latencies" | tail -1)
    min_lat=$(echo "$latencies" | head -1)
fi

echo ""
echo "=== Results ==="
echo "Total:     $TOTAL"
echo "Success:   $SUCCESSES ($((SUCCESSES * 100 / TOTAL))%)"
echo "Failures:  $FAILURES"
echo "OOM:       $OOM_COUNT"
echo ""
echo "=== Latency (ms) ==="
echo "Min:  ${min_lat:-N/A}"
echo "Avg:  ${avg:-N/A}"
echo "p50:  ${p50:-N/A}"
echo "p95:  ${p95:-N/A}"
echo "p99:  ${p99:-N/A}"
echo "Max:  ${max_lat:-N/A}"

# Cleanup
rm -rf "$RESULTS_DIR"

if [[ "$FAILURES" -gt 0 ]]; then
    echo ""
    echo "⚠ $FAILURES failures detected"
    exit 1
fi

echo ""
echo "✓ All expressions evaluated successfully"