#!/bin/bash

# ================================================================================================
# Ensemble Database Performance Benchmark
# Tests write throughput, query latency, and compression efficiency
# Target: >1000 inserts/sec, P99 <100ms, compression >5x
# ================================================================================================

set -euo pipefail

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Database connection
DB_URL="postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt"

# Test parameters
WRITE_TEST_DURATION=10  # seconds
WRITE_TEST_BATCH_SIZE=100
TARGET_WRITES_PER_SEC=1000
TARGET_P99_LATENCY_MS=100
TARGET_COMPRESSION_RATIO=5.0

echo -e "${BLUE}================================================================================================${NC}"
echo -e "${BLUE}ENSEMBLE DATABASE PERFORMANCE BENCHMARK${NC}"
echo -e "${BLUE}================================================================================================${NC}"
echo ""

# ================================================================================================
# PART 1: PRE-TEST SETUP
# ================================================================================================

echo -e "${YELLOW}[1/6] Pre-test Setup${NC}"
echo "Applying migration 023..."

psql "$DB_URL" -f /home/jgrusewski/Work/foxhunt/migrations/023_ensemble_performance_tuning.sql > /dev/null 2>&1 || {
    echo -e "${RED}❌ Migration 023 failed${NC}"
    exit 1
}

echo -e "${GREEN}✅ Migration 023 applied successfully${NC}"
echo ""

# Enable pg_stat_statements for query monitoring
psql "$DB_URL" -c "CREATE EXTENSION IF NOT EXISTS pg_stat_statements;" > /dev/null 2>&1
psql "$DB_URL" -c "SELECT pg_stat_statements_reset();" > /dev/null 2>&1

echo -e "${GREEN}✅ pg_stat_statements enabled and reset${NC}"
echo ""

# ================================================================================================
# PART 2: WRITE THROUGHPUT TEST
# ================================================================================================

echo -e "${YELLOW}[2/6] Write Throughput Test (${WRITE_TEST_DURATION} seconds)${NC}"
echo "Target: ${TARGET_WRITES_PER_SEC} inserts/sec"
echo ""

# Generate test data and insert in batches
START_TIME=$(date +%s)
TOTAL_INSERTS=0

for i in $(seq 1 $WRITE_TEST_DURATION); do
    BATCH_JSON=$(cat <<EOF
[
$(for j in $(seq 1 $WRITE_TEST_BATCH_SIZE); do
    TIMESTAMP=$(date -u -Iseconds)
    SYMBOL=$(printf "SYM%02d" $((RANDOM % 10 + 1)))
    ACTION=$(printf "%s" "$(shuf -n1 -e BUY SELL HOLD)")
    cat <<INNER_EOF
{
    "timestamp": "$TIMESTAMP",
    "symbol": "$SYMBOL",
    "ensemble_action": "$ACTION",
    "ensemble_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "ensemble_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "disagreement_rate": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "dqn_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "dqn_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "dqn_weight": 0.25,
    "dqn_vote": "$ACTION",
    "ppo_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "ppo_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "ppo_weight": 0.25,
    "ppo_vote": "$ACTION",
    "mamba2_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "mamba2_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "mamba2_weight": 0.25,
    "mamba2_vote": "$ACTION",
    "tft_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "tft_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "tft_weight": 0.25,
    "tft_vote": "$ACTION",
    "inference_latency_us": $((RANDOM % 10000 + 1000)),
    "aggregation_latency_us": $((RANDOM % 1000 + 100))
}$(if [ $j -lt $WRITE_TEST_BATCH_SIZE ]; then echo ","; fi)
INNER_EOF
done)
]
EOF
)

    # Insert batch using bulk function
    BATCH_START=$(date +%s%N)
    ROWS_INSERTED=$(psql "$DB_URL" -t -c "SELECT insert_ensemble_predictions_bulk('$BATCH_JSON'::JSONB);" 2>/dev/null | tr -d ' ')
    BATCH_END=$(date +%s%N)
    BATCH_DURATION_MS=$(( (BATCH_END - BATCH_START) / 1000000 ))

    TOTAL_INSERTS=$((TOTAL_INSERTS + ROWS_INSERTED))
    echo -ne "\rBatch $i: ${ROWS_INSERTED} rows in ${BATCH_DURATION_MS}ms | Total: ${TOTAL_INSERTS} rows"
done

END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
WRITES_PER_SEC=$((TOTAL_INSERTS / DURATION))

echo ""
echo ""
echo -e "Total inserts: ${BLUE}${TOTAL_INSERTS}${NC}"
echo -e "Duration: ${BLUE}${DURATION}${NC} seconds"
echo -e "Write throughput: ${BLUE}${WRITES_PER_SEC}${NC} inserts/sec"

if [ $WRITES_PER_SEC -ge $TARGET_WRITES_PER_SEC ]; then
    echo -e "${GREEN}✅ PASS: Write throughput ${WRITES_PER_SEC}/sec >= target ${TARGET_WRITES_PER_SEC}/sec${NC}"
else
    echo -e "${RED}❌ FAIL: Write throughput ${WRITES_PER_SEC}/sec < target ${TARGET_WRITES_PER_SEC}/sec${NC}"
fi
echo ""

# ================================================================================================
# PART 3: QUERY LATENCY BENCHMARKS (26 production queries)
# ================================================================================================

echo -e "${YELLOW}[3/6] Query Latency Benchmarks (26 production queries)${NC}"
echo "Target: P99 < ${TARGET_P99_LATENCY_MS}ms"
echo ""

# Array of query names and SQL
declare -a QUERIES=(
    "Q1: Recent predictions by symbol|SELECT * FROM ensemble_predictions WHERE symbol = 'SYM01' ORDER BY timestamp DESC LIMIT 100"
    "Q2: High disagreement events|SELECT * FROM ensemble_predictions WHERE disagreement_rate > 0.5 ORDER BY timestamp DESC LIMIT 100"
    "Q3: P&L attribution by symbol|SELECT symbol, SUM(pnl) as total_pnl FROM ensemble_predictions WHERE pnl IS NOT NULL GROUP BY symbol"
    "Q4: Model performance by symbol|SELECT model_id, symbol, AVG(accuracy) FROM model_performance_attribution WHERE window_hours = 24 GROUP BY model_id, symbol"
    "Q5: Top performers 24h|SELECT * FROM get_top_models_24h('SYM01', 5)"
    "Q6: Ensemble hourly metrics|SELECT * FROM ensemble_performance_hourly WHERE symbol = 'SYM01' ORDER BY bucket DESC LIMIT 48"
    "Q7: Model correlation 7d|SELECT * FROM calculate_model_correlation_7d('SYM01')"
    "Q8: High disagreement 24h|SELECT * FROM get_high_disagreement_events_24h('SYM01', 0.5, 100)"
    "Q9: Write throughput 5min|SELECT * FROM ensemble_write_throughput_5min"
    "Q10: Action distribution|SELECT ensemble_action, COUNT(*) FROM ensemble_predictions GROUP BY ensemble_action"
    "Q11: Avg confidence by action|SELECT ensemble_action, AVG(ensemble_confidence) FROM ensemble_predictions GROUP BY ensemble_action"
    "Q12: Latency P99|SELECT PERCENTILE_CONT(0.99) WITHIN GROUP (ORDER BY inference_latency_us) FROM ensemble_predictions"
    "Q13: Model vote agreement|SELECT COUNT(*) FROM ensemble_predictions WHERE dqn_vote = ppo_vote AND ppo_vote = mamba2_vote AND mamba2_vote = tft_vote"
    "Q14: Recent orders with P&L|SELECT * FROM ensemble_predictions WHERE order_id IS NOT NULL ORDER BY timestamp DESC LIMIT 100"
    "Q15: Win rate by symbol|SELECT symbol, COUNT(CASE WHEN pnl > 0 THEN 1 END)::FLOAT / NULLIF(COUNT(*), 0) FROM ensemble_predictions WHERE pnl IS NOT NULL GROUP BY symbol"
    "Q16: Model performance hourly|SELECT * FROM model_performance_hourly WHERE model_id = 'DQN' ORDER BY bucket DESC LIMIT 24"
    "Q17: Ensemble weekly summary|SELECT * FROM ensemble_performance_weekly ORDER BY bucket DESC LIMIT 12"
    "Q18: Avg Sharpe by model|SELECT model_id, AVG(sharpe_ratio) FROM model_performance_attribution WHERE window_hours = 24 GROUP BY model_id"
    "Q19: Max drawdown by symbol|SELECT symbol, MAX(max_drawdown) FROM model_performance_attribution WHERE window_hours = 168 GROUP BY symbol"
    "Q20: Checkpoint performance|SELECT dqn_checkpoint_id, AVG(ensemble_confidence) FROM ensemble_predictions WHERE dqn_checkpoint_id IS NOT NULL GROUP BY dqn_checkpoint_id"
    "Q21: Time-weighted avg signal|SELECT time_bucket('1 hour', timestamp), AVG(ensemble_signal) FROM ensemble_predictions GROUP BY 1 ORDER BY 1 DESC LIMIT 24"
    "Q22: Disagreement rate trend|SELECT time_bucket('1 day', timestamp), AVG(disagreement_rate) FROM ensemble_predictions GROUP BY 1 ORDER BY 1 DESC LIMIT 30"
    "Q23: Model weight distribution|SELECT model_id, AVG(avg_weight) FROM model_performance_attribution WHERE window_hours = 1 GROUP BY model_id"
    "Q24: Recent high confidence|SELECT * FROM ensemble_predictions WHERE ensemble_confidence > 0.8 ORDER BY timestamp DESC LIMIT 100"
    "Q25: P&L by action type|SELECT ensemble_action, SUM(pnl) FROM ensemble_predictions WHERE pnl IS NOT NULL GROUP BY ensemble_action"
    "Q26: Inference latency trend|SELECT time_bucket('1 hour', timestamp), AVG(inference_latency_us), PERCENTILE_CONT(0.99) WITHIN GROUP (ORDER BY inference_latency_us) FROM ensemble_predictions GROUP BY 1 ORDER BY 1 DESC LIMIT 24"
)

# Run each query 5 times and collect timing
QUERY_COUNT=${#QUERIES[@]}
declare -a QUERY_TIMES=()

for i in "${!QUERIES[@]}"; do
    IFS='|' read -r QUERY_NAME QUERY_SQL <<< "${QUERIES[$i]}"

    # Run query 5 times
    TIMES=()
    for run in {1..5}; do
        START=$(date +%s%N)
        psql "$DB_URL" -c "$QUERY_SQL" > /dev/null 2>&1
        END=$(date +%s%N)
        DURATION_MS=$(( (END - START) / 1000000 ))
        TIMES+=($DURATION_MS)
    done

    # Calculate median time
    IFS=$'\n' SORTED_TIMES=($(sort -n <<<"${TIMES[*]}"))
    MEDIAN_TIME=${SORTED_TIMES[2]}
    QUERY_TIMES+=($MEDIAN_TIME)

    echo -e "${QUERY_NAME}: ${BLUE}${MEDIAN_TIME}ms${NC}"
done

echo ""

# Calculate P99 latency
IFS=$'\n' SORTED_QUERY_TIMES=($(sort -n <<<"${QUERY_TIMES[*]}"))
P99_INDEX=$(( (QUERY_COUNT * 99) / 100 ))
P99_LATENCY=${SORTED_QUERY_TIMES[$P99_INDEX]}

echo -e "P99 Query Latency: ${BLUE}${P99_LATENCY}ms${NC}"

if [ $P99_LATENCY -le $TARGET_P99_LATENCY_MS ]; then
    echo -e "${GREEN}✅ PASS: P99 latency ${P99_LATENCY}ms <= target ${TARGET_P99_LATENCY_MS}ms${NC}"
else
    echo -e "${RED}❌ FAIL: P99 latency ${P99_LATENCY}ms > target ${TARGET_P99_LATENCY_MS}ms${NC}"
fi
echo ""

# ================================================================================================
# PART 4: COMPRESSION RATIO TEST
# ================================================================================================

echo -e "${YELLOW}[4/6] Compression Ratio Test${NC}"
echo "Target: Compression ratio > ${TARGET_COMPRESSION_RATIO}x"
echo ""

# Insert old data (8 days ago) to trigger compression
echo "Inserting old data for compression test..."

OLD_DATA_JSON=$(cat <<EOF
[
$(for j in $(seq 1 1000); do
    OLD_TIMESTAMP=$(date -u -Iseconds -d '8 days ago')
    SYMBOL=$(printf "SYM%02d" $((RANDOM % 10 + 1)))
    ACTION=$(printf "%s" "$(shuf -n1 -e BUY SELL HOLD)")
    cat <<INNER_EOF
{
    "timestamp": "$OLD_TIMESTAMP",
    "symbol": "$SYMBOL",
    "ensemble_action": "$ACTION",
    "ensemble_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "ensemble_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "disagreement_rate": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "dqn_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "dqn_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "dqn_weight": 0.25,
    "dqn_vote": "$ACTION",
    "ppo_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "ppo_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "ppo_weight": 0.25,
    "ppo_vote": "$ACTION",
    "mamba2_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "mamba2_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "mamba2_weight": 0.25,
    "mamba2_vote": "$ACTION",
    "tft_signal": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()*2-1}'),
    "tft_confidence": $(awk -v seed=$RANDOM 'BEGIN{srand(seed); print rand()}'),
    "tft_weight": 0.25,
    "tft_vote": "$ACTION",
    "inference_latency_us": $((RANDOM % 10000 + 1000)),
    "aggregation_latency_us": $((RANDOM % 1000 + 100))
}$(if [ $j -lt 1000 ]; then echo ","; fi)
INNER_EOF
done)
]
EOF
)

psql "$DB_URL" -t -c "SELECT insert_ensemble_predictions_bulk('$OLD_DATA_JSON'::JSONB);" > /dev/null 2>&1

echo "Triggering manual compression..."
psql "$DB_URL" -c "SELECT compress_chunk(i.chunk_schema || '.' || i.chunk_name) FROM timescaledb_information.chunks i WHERE i.hypertable_name = 'ensemble_predictions' AND i.is_compressed = false AND i.range_start < NOW() - INTERVAL '7 days';" > /dev/null 2>&1

# Wait for compression to complete
sleep 2

# Check compression ratio
COMPRESSION_STATS=$(psql "$DB_URL" -t -c "SELECT AVG(before_compression_total_bytes::FLOAT / NULLIF(after_compression_total_bytes, 0)) FROM timescaledb_information.compressed_chunk_stats WHERE hypertable_name = 'ensemble_predictions';" | tr -d ' ')

if [ -z "$COMPRESSION_STATS" ] || [ "$COMPRESSION_STATS" == "" ]; then
    echo -e "${YELLOW}⚠️  No compressed chunks yet (data too recent)${NC}"
    echo -e "${BLUE}Note: Compression will trigger automatically after 7 days${NC}"
else
    COMPRESSION_RATIO=$(printf "%.2f" "$COMPRESSION_STATS")
    echo -e "Compression ratio: ${BLUE}${COMPRESSION_RATIO}x${NC}"

    if (( $(echo "$COMPRESSION_RATIO >= $TARGET_COMPRESSION_RATIO" | bc -l) )); then
        echo -e "${GREEN}✅ PASS: Compression ratio ${COMPRESSION_RATIO}x >= target ${TARGET_COMPRESSION_RATIO}x${NC}"
    else
        echo -e "${RED}❌ FAIL: Compression ratio ${COMPRESSION_RATIO}x < target ${TARGET_COMPRESSION_RATIO}x${NC}"
    fi
fi
echo ""

# ================================================================================================
# PART 5: INDEX EFFICIENCY TEST
# ================================================================================================

echo -e "${YELLOW}[5/6] Index Efficiency Test${NC}"
echo ""

# Check index usage statistics
psql "$DB_URL" -c "
SELECT
    schemaname,
    tablename,
    indexname,
    idx_scan as index_scans,
    idx_tup_read as tuples_read,
    idx_tup_fetch as tuples_fetched,
    pg_size_pretty(pg_relation_size(indexrelid)) as index_size
FROM pg_stat_user_indexes
WHERE tablename IN ('ensemble_predictions', 'model_performance_attribution')
ORDER BY idx_scan DESC, tablename;
"

echo ""

# ================================================================================================
# PART 6: CONTINUOUS AGGREGATE TEST
# ================================================================================================

echo -e "${YELLOW}[6/6] Continuous Aggregate Refresh Test${NC}"
echo ""

# Manually refresh continuous aggregates
echo "Refreshing continuous aggregates..."

psql "$DB_URL" -c "CALL refresh_continuous_aggregate('ensemble_performance_5min', NOW() - INTERVAL '1 hour', NOW());" > /dev/null 2>&1
psql "$DB_URL" -c "CALL refresh_continuous_aggregate('model_performance_hourly', NOW() - INTERVAL '6 hours', NOW());" > /dev/null 2>&1
psql "$DB_URL" -c "CALL refresh_continuous_aggregate('ensemble_performance_weekly', NOW() - INTERVAL '1 week', NOW());" > /dev/null 2>&1

echo -e "${GREEN}✅ Continuous aggregates refreshed${NC}"
echo ""

# Check continuous aggregate sizes
psql "$DB_URL" -c "
SELECT
    view_name,
    pg_size_pretty(pg_total_relation_size(format('%I.%I', view_schema, view_name)::regclass)) as total_size
FROM timescaledb_information.continuous_aggregates
WHERE view_name IN ('ensemble_performance_5min', 'model_performance_hourly', 'ensemble_performance_weekly')
ORDER BY view_name;
"

echo ""

# ================================================================================================
# FINAL SUMMARY
# ================================================================================================

echo -e "${BLUE}================================================================================================${NC}"
echo -e "${BLUE}BENCHMARK SUMMARY${NC}"
echo -e "${BLUE}================================================================================================${NC}"
echo ""

echo -e "Write Throughput:     ${BLUE}${WRITES_PER_SEC}/sec${NC} (target: ${TARGET_WRITES_PER_SEC}/sec)"
echo -e "P99 Query Latency:    ${BLUE}${P99_LATENCY}ms${NC} (target: <${TARGET_P99_LATENCY_MS}ms)"
if [ -z "$COMPRESSION_STATS" ] || [ "$COMPRESSION_STATS" == "" ]; then
    echo -e "Compression Ratio:    ${YELLOW}N/A (data too recent)${NC} (target: >${TARGET_COMPRESSION_RATIO}x)"
else
    echo -e "Compression Ratio:    ${BLUE}${COMPRESSION_RATIO}x${NC} (target: >${TARGET_COMPRESSION_RATIO}x)"
fi

echo ""

# Overall pass/fail
PASS_COUNT=0
TOTAL_TESTS=3

if [ $WRITES_PER_SEC -ge $TARGET_WRITES_PER_SEC ]; then
    PASS_COUNT=$((PASS_COUNT + 1))
fi

if [ $P99_LATENCY -le $TARGET_P99_LATENCY_MS ]; then
    PASS_COUNT=$((PASS_COUNT + 1))
fi

if [ -n "$COMPRESSION_STATS" ] && [ "$COMPRESSION_STATS" != "" ]; then
    if (( $(echo "$COMPRESSION_RATIO >= $TARGET_COMPRESSION_RATIO" | bc -l) )); then
        PASS_COUNT=$((PASS_COUNT + 1))
    fi
else
    TOTAL_TESTS=2  # Exclude compression test if no data
fi

echo -e "${BLUE}Tests Passed: ${PASS_COUNT}/${TOTAL_TESTS}${NC}"
echo ""

if [ $PASS_COUNT -eq $TOTAL_TESTS ]; then
    echo -e "${GREEN}✅ ALL TESTS PASSED - Database optimized for production${NC}"
    exit 0
else
    echo -e "${YELLOW}⚠️  Some tests did not meet targets - review optimization strategies${NC}"
    exit 1
fi