#!/bin/bash
# Script để fix và test ngrok tunnel với PostgreSQL
# Sử dụng: ./fix_ngrok_postgres.sh

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
OPS_DIR="$PROJECT_ROOT/ops"
ENV_TUNNEL="$OPS_DIR/.env.tunnel"

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

log() {
    echo -e "${GREEN}[FIX-NGROK]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[FIX-NGROK]${NC} $1"
}

error() {
    echo -e "${RED}[FIX-NGROK]${NC} $1"
}

# Default values
POSTGRES_PORT=${POSTGRES_PORT:-5543}
POSTGRES_USER=${POSTGRES_USER:-hue_remote}
POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-huepass123}
POSTGRES_DB=${POSTGRES_DB:-hue_portal}
NGROK_REGION=${NGROK_REGION:-ap}

log "============================================================"
log "Fix và test ngrok tunnel với PostgreSQL"
log "============================================================"

# Step 1: Check PostgreSQL
log "Bước 1: Kiểm tra PostgreSQL..."
export PGPASSWORD="$POSTGRES_PASSWORD"
if psql -h localhost -p "$POSTGRES_PORT" -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c "SELECT 1;" &>/dev/null; then
    log "✅ PostgreSQL đang chạy trên localhost:$POSTGRES_PORT"
else
    error "❌ PostgreSQL không chạy trên localhost:$POSTGRES_PORT"
    exit 1
fi

# Step 2: Check PostgreSQL listen addresses
log "Bước 2: Kiểm tra PostgreSQL config..."
LISTEN_ADDR=$(psql -h localhost -p "$POSTGRES_PORT" -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c "SHOW listen_addresses;" 2>/dev/null | xargs)
if [ "$LISTEN_ADDR" = "*" ]; then
    log "✅ PostgreSQL đang listen trên tất cả interfaces (*)"
else
    warn "⚠️  PostgreSQL listen trên: $LISTEN_ADDR (cần là *)"
fi

# Step 3: Check Docker container
log "Bước 3: Kiểm tra Docker container..."
if docker ps | grep -q "tryhardemnayproject-db-1\|postgres.*5543"; then
    CONTAINER_NAME=$(docker ps | grep -E "tryhardemnayproject-db-1|postgres.*5543" | awk '{print $NF}' | head -1)
    log "✅ Tìm thấy PostgreSQL container: $CONTAINER_NAME"
    
    # Check pg_hba.conf
    log "Bước 4: Kiểm tra pg_hba.conf..."
    if docker exec "$CONTAINER_NAME" grep -q "0.0.0.0/0" /var/lib/postgresql/data/pg_hba.conf 2>/dev/null; then
        log "✅ pg_hba.conf đã có rule cho 0.0.0.0/0"
    else
        warn "⚠️  pg_hba.conf chưa có rule cho 0.0.0.0/0, đang thêm..."
        docker exec "$CONTAINER_NAME" bash -c "echo 'host all all 0.0.0.0/0 md5' >> /var/lib/postgresql/data/pg_hba.conf" 2>/dev/null || true
        log "✅ Đã thêm rule vào pg_hba.conf (cần restart container để apply)"
    fi
else
    warn "⚠️  Không tìm thấy PostgreSQL container"
fi

# Step 5: Restart ngrok
log "Bước 5: Khởi động lại ngrok tunnel..."
pkill -f "ngrok tcp" 2>/dev/null || true
sleep 2

# Start ngrok
nohup ngrok tcp "$POSTGRES_PORT" --region "$NGROK_REGION" --log=stdout > "$OPS_DIR/ngrok.log" 2>&1 &
NGROK_PID=$!

# Wait for ngrok
for i in {1..10}; do
    sleep 1
    if curl -s http://127.0.0.1:4040/api/tunnels &>/dev/null; then
        NGROK_URL=$(curl -s http://127.0.0.1:4040/api/tunnels | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    for tunnel in data.get('tunnels', []):
        if tunnel.get('proto') == 'tcp':
            print(tunnel.get('public_url', '').replace('tcp://', ''))
            sys.exit(0)
except:
    pass
" 2>/dev/null)
        if [ -n "$NGROK_URL" ]; then
            NGROK_HOST=$(echo "$NGROK_URL" | cut -d':' -f1)
            NGROK_PORT=$(echo "$NGROK_URL" | cut -d':' -f2)
            log "✅ Ngrok đã sẵn sàng: tcp://$NGROK_HOST:$NGROK_PORT"
            break
        fi
    fi
done

if [ -z "$NGROK_HOST" ] || [ -z "$NGROK_PORT" ]; then
    error "❌ Không thể khởi động ngrok"
    exit 1
fi

# Step 6: Test connection
log "Bước 6: Test kết nối PostgreSQL qua ngrok..."
sleep 3

# Test với psql
if timeout 10 psql -h "$NGROK_HOST" -p "$NGROK_PORT" -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c "SELECT 1;" &>/dev/null; then
    log "✅✅✅ KẾT NỐI THÀNH CÔNG QUA NGROK!"
    
    # Get document count
    DOC_COUNT=$(psql -h "$NGROK_HOST" -p "$NGROK_PORT" -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c "SELECT COUNT(*) FROM core_legaldocument;" 2>/dev/null | xargs)
    log "   Số lượng documents: $DOC_COUNT"
    
    # Save tunnel info
    log "Bước 7: Lưu thông tin tunnel..."
    mkdir -p "$OPS_DIR"
    cat > "$ENV_TUNNEL" << EOF
# Ngrok tunnel info - Generated by fix_ngrok_postgres.sh
# Last updated: $(date -u +"%Y-%m-%dT%H:%M:%S")

PG_TUNNEL_HOST=$NGROK_HOST
PG_TUNNEL_PORT=$NGROK_PORT
PG_TUNNEL_LOCAL_PORT=$POSTGRES_PORT
PG_TUNNEL_USER=$POSTGRES_USER
PG_TUNNEL_PASSWORD=$POSTGRES_PASSWORD
PG_TUNNEL_DB=$POSTGRES_DB
DATABASE_URL=postgres://$POSTGRES_USER:$POSTGRES_PASSWORD@$NGROK_HOST:$NGROK_PORT/$POSTGRES_DB
HF_SPACE_ID=${HF_SPACE_ID:-davidtran999/hue-portal-backend}
POSTGRES_HOST=${POSTGRES_HOST:-localhost}
POSTGRES_PORT=$POSTGRES_PORT
POSTGRES_USER=$POSTGRES_USER
POSTGRES_PASSWORD=$POSTGRES_PASSWORD
POSTGRES_DB=$POSTGRES_DB
PG_TUNNEL_LAST_UPDATED=$(date -u +"%Y-%m-%dT%H:%M:%S")
EOF
    log "✅ Đã lưu tunnel info vào $ENV_TUNNEL"
    
    # Update Hugging Face Space
    log "Bước 8: Cập nhật DATABASE_URL lên Hugging Face Space..."
    if [ -f "$PROJECT_ROOT/backend/venv/bin/python" ]; then
        "$PROJECT_ROOT/backend/venv/bin/python" "$SCRIPT_DIR/start_ngrok_and_set_db.py" 2>&1 | grep -E "(✅|❌|📌)" || true
    fi
    
    log "============================================================"
    log "✅✅✅ HOÀN TẤT!"
    log "   PostgreSQL: localhost:$POSTGRES_PORT"
    log "   Ngrok tunnel: tcp://$NGROK_HOST:$NGROK_PORT"
    log "   DATABASE_URL: postgres://$POSTGRES_USER:***@$NGROK_HOST:$NGROK_PORT/$POSTGRES_DB"
    log "============================================================"
else
    error "❌ Vẫn không kết nối được qua ngrok"
    error "   Có thể cần:"
    error "   1. Restart PostgreSQL container để apply pg_hba.conf"
    error "   2. Hoặc dùng SSH tunnel thay vì ngrok"
    error "   3. Hoặc expose PostgreSQL với network_mode: host"
    exit 1
fi