#!/usr/bin/env bash
# =============================================================================
# run-tests.sh — Full-stack integration test for ja4-platform
#
# Starts the entire pipeline in Docker Compose, generates traffic, and verifies
# data flows end-to-end: Apache → mod-reqin-log → correlator → ClickHouse
#                         sentinel ↗                    ↓
#                                              bot-detector → ML scores
#                                              dashboard API ← query
#
# Usage:
#   ./run-tests.sh              # run tests (build + up + test + down)
#   ./run-tests.sh --no-down    # keep stack running after tests (for debugging)
#   ./run-tests.sh --build-only # build images only, don't run tests
# =============================================================================
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

KEEP_UP=false
BUILD_ONLY=false
TESTS_PASSED=0
TESTS_FAILED=0

for arg in "$@"; do
    case "$arg" in
        --no-down)    KEEP_UP=true ;;
        --build-only) BUILD_ONLY=true ;;
    esac
done

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
log()  { echo -e "${CYAN}[test]${NC} $(date +%H:%M:%S) $*"; }
pass() { echo -e "${GREEN}  ✓ $*${NC}"; TESTS_PASSED=$((TESTS_PASSED + 1)); }
fail() { echo -e "${RED}  ✗ $*${NC}"; TESTS_FAILED=$((TESTS_FAILED + 1)); }
warn() { echo -e "${YELLOW}  ⚠ $*${NC}"; }

cleanup() {
    if [ "$KEEP_UP" = false ]; then
        log "Tearing down stack..."
        docker compose down -v --remove-orphans 2>/dev/null || true
    else
        log "Stack left running (--no-down). Stop with: docker compose down -v"
    fi
}
trap cleanup EXIT

ch_query() {
    docker compose exec -T clickhouse clickhouse-client --query "$1" 2>/dev/null
}

wait_for_service() {
    local service="$1"
    local max_wait="${2:-120}"
    log "Waiting for $service to be healthy (max ${max_wait}s)..."
    local elapsed=0
    while [ $elapsed -lt "$max_wait" ]; do
        local status
        status=$(docker compose ps --format json "$service" 2>/dev/null | python3 -c "
import sys, json
for line in sys.stdin:
    d = json.loads(line)
    print(d.get('Health','unknown'))
" 2>/dev/null || echo "unknown")
        if [ "$status" = "healthy" ]; then
            log "$service is healthy (${elapsed}s)"
            return 0
        fi
        sleep 2
        elapsed=$((elapsed + 2))
    done
    log "ERROR: $service not healthy after ${max_wait}s"
    docker compose logs --tail=30 "$service"
    return 1
}

# =============================================================================
# Phase 1: Build
# =============================================================================
log "============================================"
log "Phase 1: Building images"
log "============================================"
docker compose build --parallel 2>&1 | tail -20

if [ "$BUILD_ONLY" = true ]; then
    log "Build complete (--build-only). Exiting."
    exit 0
fi

# =============================================================================
# Phase 2: Start stack (always fresh — destroy volumes to reset DB)
# =============================================================================
log "============================================"
log "Phase 2: Starting stack (fresh DB)"
log "============================================"

# Always destroy volumes so ClickHouse reinitializes schema from scratch.
# This guarantees test isolation across runs.
log "Resetting state (docker compose down -v)..."
docker compose down -v --remove-orphans 2>/dev/null || true

docker compose up -d

wait_for_service clickhouse 120
wait_for_service platform 120
wait_for_service dashboard 60

# =============================================================================
# Phase 3: Verify ClickHouse schema
# =============================================================================
log "============================================"
log "Phase 3: Verifying ClickHouse schema"
log "============================================"

# Check databases exist
DB_COUNT=$(ch_query "SELECT count() FROM system.databases WHERE name IN ('ja4_logs','ja4_processing')")
if [ "$DB_COUNT" = "2" ]; then
    pass "Both databases created (ja4_logs, ja4_processing)"
else
    fail "Expected 2 databases, got $DB_COUNT"
fi

# Check key tables
for table in "ja4_logs.http_logs_raw" "ja4_logs.http_logs" "ja4_processing.ml_detected_anomalies" "ja4_processing.agg_host_ip_ja4_1h"; do
    db=$(echo "$table" | cut -d. -f1)
    tbl=$(echo "$table" | cut -d. -f2)
    EXISTS=$(ch_query "SELECT count() FROM system.tables WHERE database='$db' AND name='$tbl'")
    if [ "$EXISTS" = "1" ]; then
        pass "Table $table exists"
    else
        fail "Table $table missing"
    fi
done

# Check users
for user in data_writer analyst; do
    EXISTS=$(ch_query "SELECT count() FROM system.users WHERE name='$user'")
    if [ "$EXISTS" = "1" ]; then
        pass "User '$user' created"
    else
        fail "User '$user' missing"
    fi
done

# =============================================================================
# Phase 4: Seed ClickHouse + Generate test traffic
# =============================================================================
log "============================================"
log "Phase 4a: Seeding ClickHouse with synthetic data"
log "============================================"

# The seeder inserts directly into http_logs_raw, triggering all MVs:
#   http_logs_raw → mv_http_logs → http_logs → mv_agg_host_ip_ja4_1h → agg_host_ip_ja4_1h
# This pre-populates:
#   - ~350,000 rows from 14,000 browser IPs (ISP ASN ranges → asn_label='isp')
#   - ~100,000 rows from 3,000 scanner IPs (datacenter ASN → ML anomaly candidates)
#   -  ~30,000 rows from 2,000 legit bot IPs (from bot_ip.csv CIDRs)
#   -  ~20,000 rows from 1,000 AI bot IPs (datacenter ranges)
# After seeding, bot_detector has ≥500 human rows → can train and run.
log "Running seed_clickhouse.py..."
if docker compose exec -T traffic-gen python /app/seed_clickhouse.py \
    --host clickhouse --port 8123 --user default --password ""; then
    pass "ClickHouse seeded (500K rows: 350K browser + 100K scanner + 30K legit-bot + 20K AI-bot)"
else
    warn "Seeder reported errors (pipeline verification will show impact)"
fi

log "============================================"
log "Phase 4b: Generating live test traffic via Apache"
log "============================================"

# Live traffic crosses the Docker network so sentinel can capture TLS handshakes.
# X-Forwarded-For is always set — mod_remoteip updates r->useragent_ip → diverse src_ips.
log "Starting traffic generator (500 requests, 10 workers)..."
if docker compose exec -T traffic-gen python /app/generate_traffic.py \
    --host platform --http-port 80 --https-port 443 \
    --requests 500 --workers 10; then
    pass "Traffic generation complete (500 requests with diverse XFF IPs: browsers, bots)"
else
    warn "Traffic generator reported some errors (>80% success still passes)"
fi

# Wait for correlator to flush all batches to ClickHouse
log "Waiting 20s for correlator to flush and bot-detector first cycle..."
sleep 20

# =============================================================================
# Phase 5: Verify data pipeline
# =============================================================================
log "============================================"
log "Phase 5: Verifying data pipeline"
log "============================================"

# 5a. Raw logs ingested
RAW_COUNT=$(ch_query "SELECT count() FROM ja4_logs.http_logs_raw")
if [ "$RAW_COUNT" -gt 0 ] 2>/dev/null; then
    pass "Raw logs ingested: $RAW_COUNT rows in http_logs_raw (seeder + live traffic)"
else
    fail "No raw logs in http_logs_raw (correlator → ClickHouse failed)"
    # Debug
    log "Correlator logs:"
    docker compose logs --tail=30 platform 2>&1 | grep -i "correlator\|error\|clickhouse" | head -20
fi

# 5b. Parsed logs via materialized view
PARSED_COUNT=$(ch_query "SELECT count() FROM ja4_logs.http_logs")
if [ "$PARSED_COUNT" -gt 0 ] 2>/dev/null; then
    pass "Parsed logs: $PARSED_COUNT rows in http_logs (MV working)"
else
    warn "No parsed logs in http_logs (MV may need INSERT trigger, or dict loading failed)"
fi

# 5c. Check a sample parsed log has expected fields
if [ "$PARSED_COUNT" -gt 0 ] 2>/dev/null; then
    # Verify variety of User-Agents (browsers + bots)
    UA_TYPES=$(ch_query "SELECT count(DISTINCT header_user_agent) FROM ja4_logs.http_logs")
    if [ "$UA_TYPES" -gt 5 ] 2>/dev/null; then
        pass "Varied User-Agents: $UA_TYPES distinct UAs in logs"
    else
        warn "Low User-Agent variety: only $UA_TYPES distinct UAs"
    fi

    # Verify HTTP method variety
    METHODS=$(ch_query "SELECT groupArray(method) FROM (SELECT DISTINCT method FROM ja4_logs.http_logs ORDER BY method)")
    pass "HTTP methods captured: $METHODS"
fi

# 5d. TLS fingerprints captured (sentinel → correlator → ClickHouse)
if [ "$PARSED_COUNT" -gt 0 ] 2>/dev/null; then
    JA4_COUNT=$(ch_query "SELECT count() FROM ja4_logs.http_logs WHERE ja4 != ''")
    JA4_UNIQ=$(ch_query "SELECT count(DISTINCT ja4) FROM ja4_logs.http_logs WHERE ja4 != ''")
    JA3_COUNT=$(ch_query "SELECT count() FROM ja4_logs.http_logs WHERE ja3 != ''")
    JA3_UNIQ=$(ch_query "SELECT count(DISTINCT ja3_hash) FROM ja4_logs.http_logs WHERE ja3_hash != ''")
    TLS_VERSIONS=$(ch_query "SELECT groupArray(tls_version) FROM (SELECT DISTINCT tls_version FROM ja4_logs.http_logs WHERE tls_version != '' ORDER BY tls_version)")

    if [ "$JA4_COUNT" -gt 0 ] 2>/dev/null; then
        pass "TLS capture: $JA4_COUNT rows with JA4 ($JA4_UNIQ unique fingerprints)"
        SAMPLE=$(ch_query "SELECT ja4, tls_version FROM ja4_logs.http_logs WHERE ja4 != '' LIMIT 1 FORMAT TabSeparated")
        log "  JA4 sample: $SAMPLE"
    else
        warn "No JA4 fingerprints (sentinel may not see traffic on eth0)"
    fi
    if [ "$JA3_COUNT" -gt 0 ] 2>/dev/null; then
        pass "TLS capture: $JA3_COUNT rows with JA3 ($JA3_UNIQ unique fingerprints)"
    fi
    if [ -n "$TLS_VERSIONS" ]; then
        pass "TLS versions seen: $TLS_VERSIONS"
    fi
fi

# 5e. Check correlator log file
CORR_LINES=$(docker compose exec -T platform wc -l < /var/log/logcorrelator/correlated.log 2>/dev/null || echo 0)
if [ "$CORR_LINES" -gt 0 ] 2>/dev/null; then
    pass "Correlator file output: $CORR_LINES lines in correlated.log"
else
    warn "Correlator file output empty"
fi

# 5f. Verify seeder data reached agg table and AI features view
AGG_COUNT=$(ch_query "SELECT count() FROM ja4_processing.agg_host_ip_ja4_1h")
HUMAN_COUNT=$(ch_query "SELECT count() FROM ja4_processing.view_ai_features_1h WHERE asn_label='isp'")
BOT_LABEL_COUNT=$(ch_query "SELECT count() FROM ja4_processing.view_ai_features_1h WHERE bot_name != ''")
UNIQ_SRC_IPS=$(ch_query "SELECT count(DISTINCT src_ip) FROM ja4_processing.view_ai_features_1h")
UNIQ_JA4=$(ch_query "SELECT count(DISTINCT ja4) FROM ja4_processing.view_ai_features_1h")

if [ "$AGG_COUNT" -gt 0 ] 2>/dev/null; then
    pass "Aggregation table populated: $AGG_COUNT sessions in agg_host_ip_ja4_1h"
else
    fail "agg_host_ip_ja4_1h empty (MV chain broken)"
fi

if [ "$HUMAN_COUNT" -ge 500 ] 2>/dev/null; then
    pass "Bot-detector baseline: $HUMAN_COUNT ISP sessions (≥500 threshold met)"
elif [ "$HUMAN_COUNT" -gt 0 ] 2>/dev/null; then
    warn "ISP sessions below threshold: $HUMAN_COUNT < 500 (bot_detector will skip cycle)"
else
    fail "No ISP sessions in view_ai_features_1h (asn_reputation CSV not loaded?)"
fi

if [ "$BOT_LABEL_COUNT" -gt 0 ] 2>/dev/null; then
    pass "Known bots labeled: $BOT_LABEL_COUNT sessions with bot_name (bot_ip/bot_ja4 dicts working)"
else
    warn "No known-bot labels in view_ai_features_1h (bot_ip.csv / bot_ja4.csv empty?)"
fi

log "  Unique src_ips: $UNIQ_SRC_IPS | Unique JA4: $UNIQ_JA4"

# =============================================================================
# Phase 6: Verify dashboard API
# =============================================================================
log "============================================"
log "Phase 6: Verifying dashboard API"
log "============================================"

# Health check (dashboard has no curl, use python urllib)
HEALTH=$(docker compose exec -T dashboard python -c "
import urllib.request, json
r = urllib.request.urlopen('http://localhost:8000/health')
print(json.loads(r.read()).get('status',''))
" 2>/dev/null || echo "FAIL")
if [ "$HEALTH" = "healthy" ] || [ "$HEALTH" = "ok" ]; then
    pass "Dashboard /health returns $HEALTH"
else
    fail "Dashboard /health failed: $HEALTH"
fi

# Metrics endpoint
METRICS_STATUS=$(docker compose exec -T dashboard python -c "
import urllib.request
try:
    r = urllib.request.urlopen('http://localhost:8000/api/metrics')
    print(r.status)
except urllib.error.HTTPError as e:
    print(e.code)
except Exception:
    print(0)
" 2>/dev/null || echo "000")
if [ "$METRICS_STATUS" = "200" ] || [ "$METRICS_STATUS" = "404" ]; then
    pass "Dashboard /api/metrics responds (HTTP $METRICS_STATUS)"
else
    fail "Dashboard /api/metrics failed (HTTP $METRICS_STATUS)"
fi

# =============================================================================
# Phase 7: Verify bot-detector
# =============================================================================
log "============================================"
log "Phase 7: Verifying bot-detector"
log "============================================"

BOT_STATUS=$(docker compose ps --format json bot-detector 2>/dev/null | python3 -c "
import sys, json
for line in sys.stdin:
    d = json.loads(line)
    print(d.get('State','unknown'))
" 2>/dev/null || echo "unknown")

if [ "$BOT_STATUS" = "running" ]; then
    pass "Bot-detector is running"
else
    warn "Bot-detector state: $BOT_STATUS"
fi

# Check if bot-detector successfully ran a detection cycle (not just SKIPPED_LOW_DATA)
BD_SCORES=$(ch_query "SELECT count() FROM ja4_processing.ml_all_scores" 2>/dev/null || echo 0)
BD_ANOMALIES=$(ch_query "SELECT count() FROM ja4_processing.ml_detected_anomalies" 2>/dev/null || echo 0)
if [ "$BD_SCORES" -gt 0 ] 2>/dev/null; then
    pass "Bot-detector scored traffic: $BD_SCORES rows in ml_all_scores, $BD_ANOMALIES anomalies detected"
else
    warn "ml_all_scores is empty — bot-detector may not have completed a cycle yet"
    warn "  (check: docker compose logs bot-detector | grep -E 'CYCLE|SKIP|train')"
fi

# =============================================================================
# Phase 8: Network capture verification (sentinel)
# =============================================================================
log "============================================"
log "Phase 8: Verifying sentinel capture"
log "============================================"

SENTINEL_RUNNING=$(docker compose exec -T platform pgrep -x sentinel > /dev/null 2>&1 && echo "yes" || echo "no")
if [ "$SENTINEL_RUNNING" = "yes" ]; then
    pass "Sentinel process is running"
else
    fail "Sentinel process not found"
    docker compose logs --tail=10 platform 2>&1 | grep -i sentinel | head -5
fi

# Check sentinel log output
SENTINEL_LOG=$(docker compose exec -T platform cat /var/log/ja4sentinel/sentinel.log 2>/dev/null | head -5 || echo "")
if [ -n "$SENTINEL_LOG" ]; then
    pass "Sentinel producing log output"
else
    warn "No sentinel log file found (may be logging to stdout only)"
fi

# =============================================================================
# Summary
# =============================================================================
echo ""
log "============================================"
log "RESULTS"
log "============================================"
TOTAL=$((TESTS_PASSED + TESTS_FAILED))
echo -e "  ${GREEN}Passed: $TESTS_PASSED${NC} / $TOTAL"
if [ "$TESTS_FAILED" -gt 0 ]; then
    echo -e "  ${RED}Failed: $TESTS_FAILED${NC} / $TOTAL"
fi
echo ""

if [ "$TESTS_FAILED" -gt 0 ]; then
    log "Some tests failed. Use --no-down to keep the stack running for debugging."
    log "Debug commands:"
    log "  docker compose logs platform"
    log "  docker compose exec platform cat /var/log/logcorrelator/correlated.log"
    log "  docker compose exec clickhouse clickhouse-client -q 'SELECT * FROM ja4_logs.http_logs_raw LIMIT 5'"
    exit 1
else
    log "All tests passed!"
    exit 0
fi