#!/bin/bash
# 修复 stock.allbyai.cn 数据库连接问题
# 在服务器上执行: sudo bash /opt/stock-app/deploy/fix-db-connection.sh

DB_PASS="stock_password_2025"
APP_DIR="/opt/stock-app"

echo "=========================================="
echo "诊断并修复 PostgreSQL 数据库连接"
echo "=========================================="

# 0. 诊断：收集当前状态
echo ""
echo "[诊断] 检查 PostgreSQL 服务状态..."
systemctl status postgresql --no-pager 2>&1 | head -10
echo ""

echo "[诊断] 检查 PostgreSQL 版本和集群..."
pg_lsclusters 2>/dev/null || echo "pg_lsclusters 不可用"
echo ""

echo "[诊断] 检查 PostgreSQL 监听端口..."
ss -tlnp | grep 5432 || netstat -tlnp 2>/dev/null | grep 5432 || echo "未检测到5432端口监听"
echo ""

echo "[诊断] 检查 pg_hba.conf 配置..."
PG_HBA=$(find /etc/postgresql -name pg_hba.conf 2>/dev/null | head -1)
if [ -n "$PG_HBA" ]; then
    echo "文件位置: $PG_HBA"
    echo "--- 当前认证配置 ---"
    grep -v '^#' "$PG_HBA" | grep -v '^$'
    echo "---"
else
    echo "未找到 pg_hba.conf"
fi
echo ""

echo "[诊断] 检查 stock-app 服务环境变量..."
systemctl show stock-app --property=Environment 2>/dev/null || echo "无法读取服务配置"
echo ""

echo "[诊断] 尝试 peer 认证连接..."
sudo -u postgres psql -c "SELECT 1 as peer_auth_ok;" 2>&1 || echo "peer 认证失败"
echo ""

echo "[诊断] 检查 stock_app 数据库是否存在..."
sudo -u postgres psql -c "SELECT datname FROM pg_database WHERE datname='stock_app';" 2>&1
echo ""

# 1. 确保 PostgreSQL 服务运行
echo "=========================================="
echo "[1/6] 确保 PostgreSQL 服务运行..."
systemctl start postgresql 2>/dev/null || true
systemctl enable postgresql 2>/dev/null || true

# 2. 重置 postgres 用户密码
echo "[2/6] 重置 postgres 用户密码..."
sudo -u postgres psql -c "ALTER USER postgres PASSWORD '${DB_PASS}';" 2>/dev/null || {
    echo "尝试使用 peer 认证重置密码..."
    sudo -u postgres psql -c "ALTER USER postgres PASSWORD '${DB_PASS}';" || {
        echo "❌ 密码重置失败，尝试重启 PostgreSQL 后重试..."
        systemctl restart postgresql
        sleep 2
        sudo -u postgres psql -c "ALTER USER postgres PASSWORD '${DB_PASS}';"
    }
}
echo "✅ 密码已重置"

# 3. 确保数据库存在
echo "[3/6] 确保 stock_app 数据库存在..."
sudo -u postgres createdb stock_app 2>/dev/null || echo "数据库已存在"

# 4. 修复 pg_hba.conf 认证配置
echo "[4/6] 检查并修复 pg_hba.conf..."
if [ -n "$PG_HBA" ]; then
    if ! grep -q "host.*all.*all.*127.0.0.1/32.*md5\|host.*all.*all.*127.0.0.1/32.*scram-sha-256" "$PG_HBA"; then
        echo "添加 localhost md5 认证规则..."
        cp "$PG_HBA" "${PG_HBA}.bak.$(date +%Y%m%d%H%M%S)"
        
        # 在文件末尾前插入规则（确保在其他 host 规则之前或文件末尾）
        if ! grep -q "^host.*all.*all.*127.0.0.1/32" "$PG_HBA"; then
            echo "host    all             all             127.0.0.1/32            md5" >> "$PG_HBA"
        fi
        if ! grep -q "^host.*all.*all.*::1/128" "$PG_HBA"; then
            echo "host    all             all             ::1/128                 md5" >> "$PG_HBA"
        fi
        
        echo "✅ pg_hba.conf 已更新，重启 PostgreSQL..."
        systemctl restart postgresql
        sleep 2
    else
        echo "✅ pg_hba.conf 认证配置正常"
    fi
else
    echo "⚠️  未找到 pg_hba.conf，跳过"
fi

# 5. 测试数据库连接
echo "[5/6] 测试数据库连接..."
PGPASSWORD="${DB_PASS}" psql -h localhost -U postgres -d stock_app -c "SELECT 1 as ok;" 2>&1
if [ $? -eq 0 ]; then
    echo "✅ 数据库连接测试通过"
else
    echo "❌ 连接测试失败！"
    echo ""
    echo "尝试替代修复方案..."
    
    # 尝试将所有 host 认证改为 md5
    if [ -n "$PG_HBA" ]; then
        echo "将 scram-sha-256 改为 md5..."
        sed -i 's/scram-sha-256/md5/g' "$PG_HBA"
        systemctl restart postgresql
        sleep 2
        
        # 重新设置密码（用 md5 格式）
        sudo -u postgres psql -c "ALTER USER postgres PASSWORD '${DB_PASS}';"
        
        PGPASSWORD="${DB_PASS}" psql -h localhost -U postgres -d stock_app -c "SELECT 1 as ok;" 2>&1
        if [ $? -eq 0 ]; then
            echo "✅ 替代方案成功"
        else
            echo "❌ 仍然失败，请手动检查 PostgreSQL 日志:"
            echo "   journalctl -u postgresql -n 50"
            echo "   cat /var/log/postgresql/*.log | tail -50"
            exit 1
        fi
    fi
fi

# 6. 重启应用服务
echo "[6/6] 重启应用服务..."
systemctl restart stock-app
systemctl restart stock-data-service
sleep 3

# 验证应用是否正常
echo ""
echo "验证应用状态..."
systemctl status stock-app --no-pager | head -5
echo ""

# 用 Python 测试应用级 DB 连接
${APP_DIR}/venv/bin/python -c "
import sys
sys.path.insert(0, '${APP_DIR}')
import os
os.environ['DB_HOST'] = 'localhost'
os.environ['DB_PORT'] = '5432'
os.environ['DB_NAME'] = 'stock_app'
os.environ['DB_USER'] = 'postgres'
os.environ['DB_PASSWORD'] = '${DB_PASS}'
from db import get_db
conn = get_db()
if conn:
    cur = conn.cursor()
    cur.execute('SELECT COUNT(*) FROM users')
    count = cur.fetchone()[0]
    print(f'✅ Python 应用级 DB 连接成功！用户数: {count}')
    conn.close()
else:
    print('❌ Python 应用级 DB 连接失败')
    sys.exit(1)
" 2>&1

echo ""
echo "=========================================="
echo "修复完成！请访问 https://stock.allbyai.cn 验证"
echo "=========================================="