#!/bin/bash

# ================================================================================================
# Quick Ensemble Database Performance Benchmark
# Fast version for immediate feedback
# ================================================================================================

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

DB_URL="postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt"

echo -e "${BLUE}Quick Ensemble Database Benchmark${NC}"
echo ""

# ================================================================================================
# TEST 1: SIMPLE WRITE THROUGHPUT (1000 rows)
# ================================================================================================

echo -e "${YELLOW}[1/4] Write Throughput Test${NC}"

# Generate simple insert test
START=$(date +%s%N)

for i in {1..10}; do
    psql "$DB_URL" -c "
    INSERT INTO ensemble_predictions (
        timestamp, symbol, ensemble_action, ensemble_signal,
        ensemble_confidence, disagreement_rate
    )
    SELECT
        NOW() - (random() * INTERVAL '1 hour'),
        'TEST_SYM',
        CASE WHEN random() < 0.33 THEN 'BUY' WHEN random() < 0.66 THEN 'SELL' ELSE 'HOLD' END,
        (random() * 2 - 1)::DOUBLE PRECISION,
        random()::DOUBLE PRECISION,
        random()::DOUBLE PRECISION
    FROM generate_series(1, 100);
    " > /dev/null 2>&1
done

END=$(date +%s%N)
DURATION_MS=$(( (END - START) / 1000000 ))
WRITES_PER_SEC=$(( 1000 * 1000 / DURATION_MS ))

echo -e "Inserted 1000 rows in ${BLUE}${DURATION_MS}ms${NC}"
echo -e "Write throughput: ${BLUE}${WRITES_PER_SEC} inserts/sec${NC}"

if [ $WRITES_PER_SEC -ge 1000 ]; then
    echo -e "${GREEN}✅ PASS: Write throughput >= 1000/sec${NC}"
else
    echo -e "${YELLOW}⚠️  WARNING: Write throughput < 1000/sec${NC}"
fi
echo ""

# ================================================================================================
# TEST 2: QUERY LATENCY (10 key queries)
# ================================================================================================

echo -e "${YELLOW}[2/4] Query Latency Test${NC}"

declare -a QUERIES=(
    "Recent predictions|SELECT * FROM ensemble_predictions ORDER BY timestamp DESC LIMIT 100"
    "High disagreement|SELECT * FROM ensemble_predictions WHERE disagreement_rate > 0.5 LIMIT 100"
    "P&L by symbol|SELECT symbol, SUM(pnl) FROM ensemble_predictions WHERE pnl IS NOT NULL GROUP BY symbol"
    "Action distribution|SELECT ensemble_action, COUNT(*) FROM ensemble_predictions GROUP BY ensemble_action"
    "Avg confidence|SELECT ensemble_action, AVG(ensemble_confidence) FROM ensemble_predictions GROUP BY ensemble_action"
    "Latency P99|SELECT PERCENTILE_CONT(0.99) WITHIN GROUP (ORDER BY inference_latency_us) FROM ensemble_predictions WHERE inference_latency_us IS NOT NULL"
    "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"
    "Recent high confidence|SELECT * FROM ensemble_predictions WHERE ensemble_confidence > 0.8 ORDER BY timestamp DESC LIMIT 100"
    "Model performance|SELECT model_id, AVG(accuracy) FROM model_performance_attribution WHERE window_hours = 24 GROUP BY model_id"
    "Hourly metrics|SELECT * FROM ensemble_performance_hourly ORDER BY bucket DESC LIMIT 24"
)

declare -a LATENCIES=()

for query_spec in "${QUERIES[@]}"; do
    IFS='|' read -r NAME SQL <<< "$query_spec"

    START=$(date +%s%N)
    psql "$DB_URL" -c "$SQL" > /dev/null 2>&1
    END=$(date +%s%N)
    LATENCY_MS=$(( (END - START) / 1000000 ))

    LATENCIES+=($LATENCY_MS)
    echo -e "${NAME}: ${BLUE}${LATENCY_MS}ms${NC}"
done

# Calculate P99
IFS=$'\n' SORTED=($(sort -n <<<"${LATENCIES[*]}"))
P99_INDEX=$(( (${#LATENCIES[@]} * 99) / 100 ))
P99_LATENCY=${SORTED[$P99_INDEX]}

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

if [ $P99_LATENCY -le 100 ]; then
    echo -e "${GREEN}✅ PASS: P99 latency <= 100ms${NC}"
else
    echo -e "${YELLOW}⚠️  WARNING: P99 latency > 100ms${NC}"
fi
echo ""

# ================================================================================================
# TEST 3: INDEX USAGE
# ================================================================================================

echo -e "${YELLOW}[3/4] Index Usage Statistics${NC}"

psql "$DB_URL" -c "
SELECT
    LEFT(indexname, 40) as index_name,
    idx_scan as scans,
    pg_size_pretty(pg_relation_size(indexrelid)) as size
FROM pg_stat_user_indexes
WHERE tablename IN ('ensemble_predictions', 'model_performance_attribution')
ORDER BY idx_scan DESC
LIMIT 10;
"

echo ""

# ================================================================================================
# TEST 4: TABLE STATISTICS
# ================================================================================================

echo -e "${YELLOW}[4/4] Table Statistics${NC}"

psql "$DB_URL" -c "
SELECT
    'ensemble_predictions' as table_name,
    COUNT(*) as row_count,
    pg_size_pretty(pg_total_relation_size('ensemble_predictions')) as total_size,
    pg_size_pretty(pg_relation_size('ensemble_predictions')) as table_size,
    pg_size_pretty(pg_indexes_size('ensemble_predictions')) as indexes_size
FROM ensemble_predictions
UNION ALL
SELECT
    'model_performance_attribution',
    COUNT(*),
    pg_size_pretty(pg_total_relation_size('model_performance_attribution')),
    pg_size_pretty(pg_relation_size('model_performance_attribution')),
    pg_size_pretty(pg_indexes_size('model_performance_attribution'))
FROM model_performance_attribution;
"

echo ""

# ================================================================================================
# 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: 1000/sec)"
echo -e "P99 Query Latency:    ${BLUE}${P99_LATENCY}ms${NC} (target: <100ms)"
echo ""

if [ $WRITES_PER_SEC -ge 1000 ] && [ $P99_LATENCY -le 100 ]; then
    echo -e "${GREEN}✅ ALL TESTS PASSED${NC}"
    exit 0
else
    echo -e "${YELLOW}⚠️  Some metrics below target - database under optimization${NC}"
    exit 0
fi